{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "name": "Flocking.ipynb",
      "provenance": [],
      "collapsed_sections": [],
      "include_colab_link": true
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "GPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "view-in-github",
        "colab_type": "text"
      },
      "source": [
        "<a href=\"https://colab.research.google.com/github/google/jax-md/blob/master/notebooks/flocking.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ktEqtIdMtovy"
      },
      "source": [
        "Copyright 2020 Google LLC\n",
        "\n",
        "Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "you may not use this file except in compliance with the License.\n",
        "You may obtain a copy of the License at\n",
        "\n",
        "     https://www.apache.org/licenses/LICENSE-2.0\n",
        "\n",
        "Unless required by applicable law or agreed to in writing, software\n",
        "distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "See the License for the specific language governing permissions and\n",
        "limitations under the License."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "VH5LodOhsl45",
        "cellView": "form"
      },
      "source": [
        "#@title Imports & Utils\n",
        "\n",
        "# Imports\n",
        "\n",
        "!pip install -q git+https://www.github.com/google/jax-md\n",
        "\n",
        "import numpy as onp\n",
        "\n",
        "from jax.config import config ; config.update('jax_enable_x64', True)\n",
        "import jax.numpy as np\n",
        "from jax import random\n",
        "from jax import jit\n",
        "from jax import vmap\n",
        "from jax import lax\n",
        "vectorize = np.vectorize\n",
        "\n",
        "from functools import partial\n",
        "\n",
        "from collections import namedtuple\n",
        "import base64\n",
        "\n",
        "import IPython\n",
        "from google.colab import output\n",
        "\n",
        "import os\n",
        "\n",
        "from jax_md import space, smap, energy, minimize, quantity, simulate, partition, util\n",
        "from jax_md.util import f32\n",
        "\n",
        "# Plotting\n",
        "\n",
        "import matplotlib.pyplot as plt\n",
        "import seaborn as sns\n",
        "  \n",
        "sns.set_style(style='white')\n",
        "\n",
        "dark_color = [56 / 256] * 3\n",
        "light_color = [213 / 256] * 3\n",
        "axis_color = 'white'\n",
        "\n",
        "def format_plot(x='', y='', grid=True):  \n",
        "  ax = plt.gca()\n",
        "  \n",
        "  ax.spines['bottom'].set_color(axis_color)\n",
        "  ax.spines['top'].set_color(axis_color) \n",
        "  ax.spines['right'].set_color(axis_color)\n",
        "  ax.spines['left'].set_color(axis_color)\n",
        "  \n",
        "  ax.tick_params(axis='x', colors=axis_color)\n",
        "  ax.tick_params(axis='y', colors=axis_color)\n",
        "  ax.yaxis.label.set_color(axis_color)\n",
        "  ax.xaxis.label.set_color(axis_color)\n",
        "  ax.set_facecolor(dark_color)\n",
        "  \n",
        "  plt.grid(grid)\n",
        "  plt.xlabel(x, fontsize=20)\n",
        "  plt.ylabel(y, fontsize=20)\n",
        "  \n",
        "def finalize_plot(shape=(1, 1)):\n",
        "  plt.gcf().patch.set_facecolor(dark_color)\n",
        "  plt.gcf().set_size_inches(\n",
        "    shape[0] * 1.5 * plt.gcf().get_size_inches()[1], \n",
        "    shape[1] * 1.5 * plt.gcf().get_size_inches()[1])\n",
        "  plt.tight_layout()\n",
        "\n",
        "# Progress Bars\n",
        "\n",
        "from IPython.display import HTML, display\n",
        "import time\n",
        "\n",
        "\n",
        "def ProgressIter(iter_fun, iter_len=0):\n",
        "  if not iter_len:\n",
        "    iter_len = len(iter_fun)\n",
        "  out = display(progress(0, iter_len), display_id=True)\n",
        "  for i, it in enumerate(iter_fun):\n",
        "    yield it\n",
        "    out.update(progress(i + 1, iter_len))\n",
        "\n",
        "def progress(value, max):\n",
        "    return HTML(\"\"\"\n",
        "        <progress\n",
        "            value='{value}'\n",
        "            max='{max}',\n",
        "            style='width: 45%'\n",
        "        >\n",
        "            {value}\n",
        "        </progress>\n",
        "    \"\"\".format(value=value, max=max))\n",
        "\n",
        "normalize = lambda v: v / np.linalg.norm(v, axis=1, keepdims=True)\n",
        "\n",
        "# Rendering\n",
        "\n",
        "renderer_code = IPython.display.HTML('''\n",
        "<canvas id=\"canvas\"></canvas>\n",
        "<script>\n",
        "  Rg = null;\n",
        "  Ng = null;\n",
        "\n",
        "  var current_scene = {\n",
        "      R: null,\n",
        "      N: null,\n",
        "      is_loaded: false,\n",
        "      frame: 0,\n",
        "      frame_count: 0,\n",
        "      boid_vertex_count: 0,\n",
        "      boid_buffer: [],\n",
        "      predator_vertex_count: 0,\n",
        "      predator_buffer: [],\n",
        "      disk_vertex_count: 0,\n",
        "      disk_buffer: null,\n",
        "      box_size: 0\n",
        "  };\n",
        "\n",
        "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
        "\n",
        "  async function load_simulation() {\n",
        "    buffer_size = 400;\n",
        "    max_frame = 800;\n",
        "\n",
        "    result = await google.colab.kernel.invokeFunction(\n",
        "        'notebook.GetObstacles', [], {});\n",
        "    data = result.data['application/json'];\n",
        "\n",
        "    if(data.hasOwnProperty('Disk')) {\n",
        "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
        "    }\n",
        "\n",
        "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
        "      console.log(i);\n",
        "      result = await google.colab.kernel.invokeFunction(\n",
        "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
        "      \n",
        "      data = result.data['application/json'];\n",
        "      current_scene = put_boids(current_scene, data);\n",
        "    }\n",
        "    current_scene.is_loaded = true;\n",
        "\n",
        "    result = await google.colab.kernel.invokeFunction(\n",
        "        'notebook.GetPredators', [], {}); \n",
        "    data = result.data['application/json'];\n",
        "    if (data.hasOwnProperty('R'))\n",
        "      current_scene = put_predators(current_scene, data);\n",
        "\n",
        "    result = await google.colab.kernel.invokeFunction(\n",
        "          'notebook.GetSimulationInfo', [], {});\n",
        "    current_scene.box_size = result.data['application/json'].box_size;\n",
        "  }\n",
        "\n",
        "  function initialize_gl() {\n",
        "    const canvas = document.getElementById(\"canvas\");\n",
        "    canvas.width = 640;\n",
        "    canvas.height = 640;\n",
        "\n",
        "    const gl = canvas.getContext(\"webgl2\");\n",
        "\n",
        "    if (!gl) {\n",
        "        alert('Unable to initialize WebGL.');\n",
        "        return;\n",
        "    }\n",
        "\n",
        "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
        "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
        "    gl.enable(gl.DEPTH_TEST);\n",
        "\n",
        "    const shader_program = initialize_shader(\n",
        "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
        "    const shader = {\n",
        "      program: shader_program,\n",
        "      attribute: {\n",
        "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
        "      },\n",
        "      uniform: {\n",
        "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
        "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
        "          color: gl.getUniformLocation(shader_program, 'color'),\n",
        "      },\n",
        "    };\n",
        "    gl.useProgram(shader_program);\n",
        "\n",
        "    const half_width = 200.0;\n",
        "\n",
        "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
        "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
        "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
        "\n",
        "    return {gl: gl, shader: shader};\n",
        "  }\n",
        "\n",
        "  var loops = 0;\n",
        "\n",
        "  function update_frame() {\n",
        "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
        "\n",
        "    if (!current_scene.is_loaded) {\n",
        "      window.requestAnimationFrame(update_frame);\n",
        "      return;\n",
        "    }\n",
        "\n",
        "    var half_width = current_scene.box_size / 2.;\n",
        "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
        "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
        "\n",
        "    if (current_scene.frame >= current_scene.frame_count) {\n",
        "      if (!current_scene.is_loaded) {\n",
        "        window.requestAnimationFrame(update_frame);\n",
        "        return;\n",
        "      }\n",
        "      loops++;\n",
        "      current_scene.frame = 0;\n",
        "    }\n",
        "\n",
        "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
        "\n",
        "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
        "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
        "    gl.vertexAttribPointer(\n",
        "      shader.attribute.vertex_position,\n",
        "      2,\n",
        "      gl.FLOAT,\n",
        "      false,\n",
        "      0,\n",
        "      0\n",
        "    );\n",
        "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
        "\n",
        "    if(current_scene.predator_buffer.length > 0)  {\n",
        "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
        "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
        "      gl.vertexAttribPointer(\n",
        "        shader.attribute.vertex_position,\n",
        "        2,\n",
        "        gl.FLOAT,\n",
        "        false,\n",
        "        0,\n",
        "        0\n",
        "      );\n",
        "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
        "    }\n",
        "    \n",
        "    if(current_scene.disk_buffer) {\n",
        "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
        "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
        "      gl.vertexAttribPointer(\n",
        "        shader.attribute.vertex_position,\n",
        "        2,\n",
        "        gl.FLOAT,\n",
        "        false,\n",
        "        0,\n",
        "        0\n",
        "      );\n",
        "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
        "    }\n",
        "\n",
        "    current_scene.frame++;\n",
        "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
        "        (current_scene.frame_count == 1 && loops < 240))\n",
        "      window.requestAnimationFrame(update_frame);\n",
        "    \n",
        "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
        "      window.requestAnimationFrame(update_frame);\n",
        "  }\n",
        "\n",
        "  function put_boids(scene, boids) {\n",
        "    const R = decode(boids['R']);\n",
        "    const R_shape = boids['R_shape'];\n",
        "    const theta = decode(boids['theta']);\n",
        "    const theta_shape = boids['theta_shape'];\n",
        "\n",
        "    function index(i, b, xy) {\n",
        "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
        "    }\n",
        "\n",
        "    var steps = R_shape[0];\n",
        "    var boids = R_shape[1];\n",
        "    var dimensions = R_shape[2];\n",
        "\n",
        "    if(dimensions != 2) {\n",
        "      alert('Can only deal with two-dimensional data.')\n",
        "    }\n",
        "\n",
        "    // First flatten the data.\n",
        "    var buffer_data = new Float32Array(boids * 6);\n",
        "    var size = 8.0;\n",
        "    for (var i = 0 ; i < steps ; i++) {\n",
        "      var buffer = gl.createBuffer();\n",
        "      for (var b = 0 ; b < boids ; b++) {\n",
        "        var xi = index(i, b, 0);\n",
        "        var yi = index(i, b, 1);\n",
        "        var ti = i * boids + b;\n",
        "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
        "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
        "        buffer_data.set([\n",
        "          R[xi] + Nx, R[yi] + Ny,\n",
        "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
        "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
        "        ], b * 6);\n",
        "      }\n",
        "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
        "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
        "\n",
        "      scene.boid_buffer.push(buffer);\n",
        "    }\n",
        "    scene.boid_vertex_count = boids * 3;\n",
        "    scene.frame_count += steps;\n",
        "    return scene;\n",
        "  }\n",
        "\n",
        "  function put_predators(scene, boids) {\n",
        "    // TODO: Unify this with the put_boids function.\n",
        "    const R = decode(boids['R']);\n",
        "    const R_shape = boids['R_shape'];\n",
        "    const theta = decode(boids['theta']);\n",
        "    const theta_shape = boids['theta_shape'];\n",
        "\n",
        "    function index(i, b, xy) {\n",
        "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
        "    }\n",
        "\n",
        "    var steps = R_shape[0];\n",
        "    var boids = R_shape[1];\n",
        "    var dimensions = R_shape[2];\n",
        "\n",
        "    if(dimensions != 2) {\n",
        "      alert('Can only deal with two-dimensional data.')\n",
        "    }\n",
        "\n",
        "    // First flatten the data.\n",
        "    var buffer_data = new Float32Array(boids * 6);\n",
        "    var size = 18.0;\n",
        "    for (var i = 0 ; i < steps ; i++) {\n",
        "      var buffer = gl.createBuffer();\n",
        "      for (var b = 0 ; b < boids ; b++) {\n",
        "        var xi = index(i, b, 0);\n",
        "        var yi = index(i, b, 1);\n",
        "        var ti = theta_shape[1] * i + b;\n",
        "        var Nx = size * Math.cos(theta[ti]);\n",
        "        var Ny = size * Math.sin(theta[ti]);\n",
        "        buffer_data.set([\n",
        "          R[xi] + Nx, R[yi] + Ny,\n",
        "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
        "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
        "        ], b * 6);\n",
        "      }\n",
        "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
        "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
        "\n",
        "      scene.predator_buffer.push(buffer);\n",
        "    }\n",
        "    scene.predator_vertex_count = boids * 3;\n",
        "    return scene;\n",
        "  }\n",
        "\n",
        "  function put_obstacle_disk(scene, disk) {\n",
        "    const R = decode(disk.R);\n",
        "    const R_shape = disk.R_shape;\n",
        "    const radius = decode(disk.D);\n",
        "    const radius_shape = disk.D_shape;\n",
        "\n",
        "    const disk_count = R_shape[0];\n",
        "    const dimensions = R_shape[1];\n",
        "    if (dimensions != 2) {\n",
        "        alert('Can only handle two-dimensional data.');\n",
        "    }\n",
        "    if (radius_shape[0] != disk_count) {\n",
        "        alert('Inconsistent disk radius count found.');\n",
        "    }\n",
        "    const segments = 32;\n",
        "\n",
        "    function index(o, xy) {\n",
        "        return o * R_shape[1] + xy;\n",
        "    }\n",
        "\n",
        "    // TODO(schsam): Use index buffers here.\n",
        "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
        "    for (var i = 0 ; i < disk_count ; i++) {\n",
        "      var xi = index(i, 0);\n",
        "      var yi = index(i, 1);\n",
        "      for (var s = 0 ; s < segments ; s++) {\n",
        "        const th = 2 * s / segments * Math.PI;\n",
        "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
        "        const rad = radius[i] * 0.8;\n",
        "        buffer_data.set([\n",
        "          R[xi], R[yi],\n",
        "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
        "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
        "        ], i * segments * 6 + s * 6);\n",
        "      }\n",
        "    }\n",
        "    var buffer = gl.createBuffer();\n",
        "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
        "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
        "    scene.disk_vertex_count = disk_count * segments * 3;\n",
        "    scene.disk_buffer = buffer;\n",
        "    return scene;\n",
        "  }\n",
        "\n",
        "  // SHADER CODE\n",
        "\n",
        "  const VERTEX_SHADER_SOURCE_2D = `\n",
        "    // Vertex Shader Program.\n",
        "    attribute vec2 vertex_position;\n",
        "    \n",
        "    uniform vec2 screen_position;\n",
        "    uniform vec2 screen_size;\n",
        "\n",
        "    void main() {\n",
        "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
        "      gl_Position = vec4(v, 0.0, 1.0);\n",
        "    }\n",
        "  `;\n",
        "\n",
        "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
        "    precision mediump float;\n",
        "\n",
        "    uniform vec4 color;\n",
        "\n",
        "    void main() {\n",
        "      gl_FragColor = color;\n",
        "    }\n",
        "  `;\n",
        "\n",
        "  function initialize_shader(\n",
        "    gl, vertex_shader_source, fragment_shader_source) {\n",
        "\n",
        "    const vertex_shader = compile_shader(\n",
        "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
        "    const fragment_shader = compile_shader(\n",
        "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
        "\n",
        "    const shader_program = gl.createProgram();\n",
        "    gl.attachShader(shader_program, vertex_shader);\n",
        "    gl.attachShader(shader_program, fragment_shader);\n",
        "    gl.linkProgram(shader_program);\n",
        "\n",
        "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
        "      alert(\n",
        "        'Unable to initialize shader program: ' + \n",
        "        gl.getProgramInfoLog(shader_program)\n",
        "        );\n",
        "        return null;\n",
        "    }\n",
        "    return shader_program;\n",
        "  }\n",
        "\n",
        "  function compile_shader(gl, type, source) {\n",
        "    const shader = gl.createShader(type);\n",
        "    gl.shaderSource(shader, source);\n",
        "    gl.compileShader(shader);\n",
        "\n",
        "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
        "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
        "      gl.deleteShader(shader);\n",
        "      return null;\n",
        "    }\n",
        "\n",
        "    return shader;\n",
        "  }\n",
        "\n",
        "  // SERIALIZATION UTILITIES\n",
        "  function decode(sBase64, nBlocksSize) {\n",
        "    var chrs = atob(atob(sBase64));\n",
        "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
        "\n",
        "    for(var i = 0 ; i < chrs.length ; i++) {\n",
        "      array[i] = chrs.charCodeAt(i);\n",
        "    }\n",
        "\n",
        "    return new Float32Array(array.buffer);\n",
        "  }\n",
        "\n",
        "  // RUN CELL\n",
        "\n",
        "  load_simulation();\n",
        "  gl_and_shader = initialize_gl();\n",
        "  var gl = gl_and_shader.gl;\n",
        "  var shader = gl_and_shader.shader;\n",
        "  update_frame();\n",
        "</script>\n",
        "''')\n",
        "\n",
        "def encode(R):\n",
        "  return base64.b64encode(onp.array(R, onp.float32).tobytes())\n",
        "\n",
        "def render(box_size, states, obstacles=None, predators=None):\n",
        "  if isinstance(states, Boids):\n",
        "    R = np.reshape(states.R, (1,) + states.R.shape)\n",
        "    theta = np.reshape(states.theta, (1,) + states.theta.shape)\n",
        "  elif isinstance(states, list):\n",
        "    if all([isinstance(x, Boids) for x in states]):\n",
        "      R, theta = zip(*states)\n",
        "      R = onp.stack(R)\n",
        "      theta = onp.stack(theta)    \n",
        "  \n",
        "  if isinstance(predators, list):\n",
        "    R_predators, theta_predators, *_ = zip(*predators)\n",
        "    R_predators = onp.stack(R_predators)\n",
        "    theta_predators = onp.stack(theta_predators)\n",
        "\n",
        "  def get_boid_states(start, end):\n",
        "    R_, theta_ = R[start:end], theta[start:end]\n",
        "    return IPython.display.JSON(data={\n",
        "        \"R_shape\": R_.shape,\n",
        "        \"R\": encode(R_), \n",
        "        \"theta_shape\": theta_.shape,\n",
        "        \"theta\": encode(theta_)\n",
        "        })\n",
        "  output.register_callback('notebook.GetBoidStates', get_boid_states)\n",
        "\n",
        "  def get_obstacles():\n",
        "    if obstacles is None:\n",
        "      return IPython.display.JSON(data={})\n",
        "    else:\n",
        "      return IPython.display.JSON(data={\n",
        "          'Disk': {\n",
        "              'R': encode(obstacles.R),\n",
        "              'R_shape': obstacles.R.shape,\n",
        "              'D': encode(obstacles.D),\n",
        "              'D_shape': obstacles.D.shape\n",
        "          }\n",
        "      })\n",
        "  output.register_callback('notebook.GetObstacles', get_obstacles)\n",
        "\n",
        "  def get_predators():\n",
        "    if predators is None:\n",
        "      return IPython.display.JSON(data={})\n",
        "    else:\n",
        "      return IPython.display.JSON(data={\n",
        "          'R': encode(R_predators),\n",
        "          'R_shape': R_predators.shape,\n",
        "          'theta': encode(theta_predators),\n",
        "          'theta_shape': theta_predators.shape\n",
        "      })\n",
        "  output.register_callback('notebook.GetPredators', get_predators)\n",
        "\n",
        "  def get_simulation_info():\n",
        "    return IPython.display.JSON(data={\n",
        "        'frames': R.shape[0],\n",
        "        'box_size': box_size\n",
        "        })\n",
        "  output.register_callback('notebook.GetSimulationInfo', get_simulation_info)\n",
        "\n",
        "  return renderer_code"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LIP8nr6ouCgF"
      },
      "source": [
        "#### **Warning**: After running the simulations in this notebook, you have to wait a moment (5 - 30 seconds) for rendering."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wxRVUJUXNb30"
      },
      "source": [
        "# Flocks, Herds, and Schools: A Distributed Behavioral Model\n",
        "\n",
        "We will go over the paper, [\"Flocks, Herds, and Schools: A Distributed Behavioral Model\"]((https://citeseerx.ist.psu.edu/viewdoc/download;jsessionid=E252054B1C02D387E8C20827CB414543?doi=10.1.1.103.7187&rep=rep1&type=pdf)) published by C. W. Reynolds in SIGGRAPH 1987. The paper itself is fantastic and, as far as a description of flocking is concerned, there is little that we can offer. Therefore, rather than go through the paper directly, we will use [JAX](https://www.github.com/google/jax) and [JAX, MD](https://www.github.com/google/jax-md) to interactively build a simulation similar to Reynolds' in colab. To simplify our discussion, we will build a two-dimensional version of Reynolds' simulation.\n",
        "\n",
        "In nature there are many examples in which large numbers of animals exhibit complex collective motion (schools of fish, flocks of birds, herds of horses, colonies of ants). In his seminal paper, Reynolds introduces a model of such collective behavior (henceforth refered to as \"flocking\") based on simple rules that can be computed locally for each entity (referred to as a \"boid\") in the flock based on its environment. This paper is written in the context of computer graphics and so Reynolds is going for biologically inspired simulations that look right rather than accuracy in any statistical sense. Ultimately, Reynolds measures success in terms of \"delight\" people find in watching the simulations; we will use a similar metric here.\n",
        "\n",
        "Note, we recommend running this notebook in \"Dark\" mode."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "oKHNcxKjfitA"
      },
      "source": [
        "## Boids\n",
        "\n",
        "Reynolds is interested in simulating bird-like entities that are described by a position, $R$, and an orientation, $\\theta$. This state can optionally augmented with extra information (for example, hunger or fear). We can define a Boids type that stores data for a collection of boids as two arrays. `R` is an `ndarray` of shape `[boid_count, spatial_dimension]` and `theta` is an ndarray of shape `[boid_count]`. An individual boid is an index into these arrays. It will often be useful to refer to the vector orientation of the boid $N = (\\cos\\theta, \\sin\\theta)$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "QO3ammBAwf4I"
      },
      "source": [
        "Boids = namedtuple('Boids', ['R', 'theta'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "he1I-LNJ0FiJ"
      },
      "source": [
        "We can instantiate a collection of boids randomly in a box of side length $L$. We will use [periodic boundary conditions](https://en.wikipedia.org/wiki/Periodic_boundary_conditions) for our simulation which means that boids will be able to wrap around the sides of the box. To do this we will use the `space.periodic` command in [JAX, MD](https://github.com/google/jax-md#spaces-spacepy)."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "bMF7YMkG1cGn"
      },
      "source": [
        "# Simulation Parameters:\n",
        "box_size = 800.0  # A float specifying the side-length of the box.\n",
        "boid_count = 200  # An integer specifying the number of boids.\n",
        "dim = 2  # The spatial dimension in which we are simulating.\n",
        "\n",
        "# Create RNG state to draw random numbers (see LINK).\n",
        "rng = random.PRNGKey(0)\n",
        "\n",
        "# Define periodic boundary conditions.\n",
        "displacement, shift = space.periodic(box_size)\n",
        "\n",
        "# Initialize the boids.\n",
        "rng, R_rng, theta_rng = random.split(rng, 3)\n",
        "\n",
        "boids = Boids(\n",
        "  R = box_size * random.uniform(R_rng, (boid_count, dim)),\n",
        "  theta = random.uniform(theta_rng, (boid_count,), maxval=2. * np.pi)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "90sBvMIubPvr",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 661
        },
        "outputId": "cc6f2a2c-cbd4-4d98-a37f-eaa50476e1da"
      },
      "source": [
        "display(render(box_size, boids))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1IGKTl_9awc_"
      },
      "source": [
        "## Dynamics\n",
        "\n",
        "Now that we have defined our boids, we have to imbue them with some rules governing their motion. Reynolds notes that in nature flocks do not seem to have a maximum size, but instead can keep acquiring new boids and grow without bound. He also comments that each boid cannot possibly be keeping track of the entire flock and must, instead, be focusing on its local neighborhood. Reynolds then proposes three simple, local, rules that boids might try to follow:\n",
        "\n",
        "\n",
        "1.   **Alignment:** Boids will try to align themselves in the direction of their neighbors.\n",
        "2.   **Avoidance:** Boids will avoid colliding with their neighbors.\n",
        "3.   **Cohesion:** Boids will try to move towards the center of mass of their neighbors.\n",
        "\n",
        "In his exposition, Reynolds is vague about the details for each of these rules and so we will take some creative liberties. We will try to phrase this problem as an energy model, so our goal will be to write down an \"energy\" function (similar to a \"loss\") $E(R, \\theta)$ such that low-energy configurations of boids satisfy each of the three rules above. \n",
        "\n",
        "\\\n",
        "\n",
        "We will write the total energy as a sum of three terms, one for each of the rules above:\n",
        "\n",
        "$$E(R, \\theta) = E_{\\text{Align}}(R, \\theta) + E_{\\text{Avoid}}(R, \\theta) + E_{\\text{Cohesion}}(R,\\theta)$$\n",
        "\n",
        "We will go through each of these rules separately below starting with alignment. Of course, any of these terms could be replaced by a learned solution. \n",
        "\n",
        "\\\n",
        "\n",
        "Once we have an energy defined in this way, configurations of boids that move along low energy trajectories might display behavior that looks appealing. However, we still have a lot of freedom to decide how we want to define dynamics over the boids. Reynolds says he uses overdamped dynamics and so we will do something similar. In particular, we will update the position of the boids so that they try to move to minimize their energy. Simultaneously , we assume that the boids are swimming (or flying / walking). We choose a particularly simple model of this to start with and assume that the boids move at a fixed speed, $v$, along whatever direction they are pointing. We will use simple forward-Euler integration. This gives an update step,\n",
        "\n",
        "$${R_i}' = R_i + \\delta t(v\\hat N_i - \\nabla_{R_i}E(R, \\theta))$$\n",
        "\n",
        "where $\\delta t$ is a timestep that we are allowed to choose. We will often refer to the force, $F^{R_i} = -\\nabla_{R_i} E(R, \\hat N)$ as the negative gradient of the energy with respect to the position of the $i$'th boid.\n",
        "\n",
        "\\\n",
        "\n",
        "We will update the orientations of the boids turn them towards \"low energy\" directions. To do this we will once again use a simple forward-Euler scheme,\n",
        "\n",
        "$$\n",
        "\\theta'_i = \\theta_i - \\delta t\\nabla_{\\theta_i}E(R,\\theta) \n",
        "$$\n",
        "\n",
        "This is just one choice of dynamics, but there are probably many that would work equally well! Feel free to play around with it. One easy improvement that one could imagine making would be to use a more sophisticated integrator. We include a Runge-Kutta 4 integrator at the top of the notebook for an adventurous reader.\n",
        "\n",
        "\\\n",
        "\n",
        "To see what this looks like before we define any interactions, we can run a simulation with $E(R,\\theta) = 0$ by first defining an `update` function that takes a boids state to a new boids state."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ftriIFKQEZ_m"
      },
      "source": [
        "@vmap\n",
        "def normal(theta):\n",
        "  return np.array([np.cos(theta), np.sin(theta)])\n",
        "\n",
        "def dynamics(energy_fn, dt, speed):\n",
        "  @jit\n",
        "  def update(_, state):\n",
        "    R, theta = state['boids']\n",
        "\n",
        "    dstate = quantity.force(energy_fn)(state)\n",
        "    dR, dtheta = dstate['boids']\n",
        "    n = normal(state['boids'].theta)\n",
        "\n",
        "    state['boids'] = Boids(shift(R, dt * (speed * n + dR)), \n",
        "                           theta + dt * dtheta)\n",
        "\n",
        "    return state\n",
        "\n",
        "  return update"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "5DH4nwjE7mKF"
      },
      "source": [
        "Now we can run a simulation and save the boid positions to a `boids_buffer` which will just be a list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "FP2D4OfC7RiO",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "074cf3d9-9ff2-4df7-f009-b420f5ada36e"
      },
      "source": [
        "update = dynamics(energy_fn=lambda state: 0., dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size, boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "UNGJJjoTAE0T"
      },
      "source": [
        "### Alignment\n",
        "\n",
        "While the above simulation works and our boids are moving happily along, it is not terribly interesting. The first thing that we can add to this simulation is the alignment rule. When writing down these rules, it is often easier to express them for a single pair of boids and then use JAX's [automatic vectorization](https://github.com/google/jax#auto-vectorization-with-vmap) via `vmap` to extend them to our entire simulation.\n",
        "\n",
        "Given a pair of boids $i$ and $j$ we would like to choose an energy function that is minimized when they are pointing in the same direction. As discussed above, one of Reynolds' requirements was locality: boids should only interact with nearby boids. To do this, we introduce a cutoff $D_{\\text{Align}}$ and ignore pairs of boids such that $\\|\\Delta R_{ij}\\| > D_{\\text{Align}}$ where $\\Delta R_{ij} = R_i - R_j$. To make it so boids react smoothly we will have the energy start out at zero when $\\|R_i - R_j\\| = D_{\\text{Align}}$ and increase smoothly as they get closer. Together, these simple ideas lead us to the following proposal,\n",
        "\n",
        "$$\\epsilon_{\\text{Align}}(\\Delta R_{ij}, \\hat N_i, \\hat N_j) = \\begin{cases}\\frac{J_{\\text{Align}}}\\alpha\\left (1 - \\frac{\\|\\Delta R_{ij}\\|}{D_{\\text{Align}}}\\right)^\\alpha(1 - \\hat N_1 \\cdot \\hat N_2)^2 & \\text{if $\\|\\Delta R_{ij}\\| < D$}\\\\ 0 & \\text{otherwise}\\end{cases}$$\n",
        "\n",
        "This energy will be maximized when $N_1$ and $N_2$ are anti-aligned and minimized when $N_1 = N_2$. In general, we would like our boids to turn to align themselves with their neighbors rather than shift their centers to move apart. Therefore, we'll insert a stop-gradient into the displacement. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4iJXkja3X1lr"
      },
      "source": [
        "def align_fn(dR, N_1, N_2, J_align, D_align, alpha):\n",
        "  dR = lax.stop_gradient(dR)\n",
        "  dr = space.distance(dR) / D_align\n",
        "  energy = J_align / alpha * (1. - dr) ** alpha * (1 - np.dot(N_1, N_2)) ** 2\n",
        "  return np.where(dr < 1.0, energy, 0.)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "SJEUu5j_Ysfp"
      },
      "source": [
        "We can plot the energy for different alignments as well as different distances between boids. We see that the energy goes to zero for large distances and when the boids are aligned."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "2IkDG8jXYwvj",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 441
        },
        "outputId": "b3594b33-1b47-4068-ba0b-9c796661a4f1"
      },
      "source": [
        "#@title Alignment Energy\n",
        "N_1 = np.array([1.0, 0.0])\n",
        "angles = np.linspace(0, np.pi, 60)\n",
        "N_2 = vmap(lambda theta: np.array([np.cos(theta), np.sin(theta)]))(angles)\n",
        "distances = np.linspace(0, 1, 5)\n",
        "dRs = vmap(lambda r: np.array([r, 0.]))(distances)\n",
        "\n",
        "fn = partial(align_fn, J_align=1., D_align=1., alpha=2.)\n",
        "energy = vmap(vmap(fn, (None, None, 0)), (0, None, None))(dRs, N_1, N_2)\n",
        "\n",
        "for d, e in zip(distances, energy):\n",
        "  plt.plot(angles, e, label='r = {}'.format(d), linewidth=3)\n",
        "\n",
        "plt.xlim([0, np.pi])\n",
        "format_plot('$\\\\theta$', '$E(r, \\\\theta)$')\n",
        "plt.legend()\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAagAAAGoCAYAAAATsnHAAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nOzdd3wUdf4/8NfMtuxmd9PYNFIg9IRmRJqA2DgRFcuJd8pxfLGXAxSNBZWv5UT5gl9Qf5bTC5bzvt6JIhxYACmCICI1hNBCKim7qbvJZtvM/P5YmCSkJ7s7W97Px4OH+czO7r4nwbyY2c983szEiRMFEEIIIX6GlboAQgghpD0UUIQQQvwSBRQhhBC/RAFFCCHEL1FAEUII8UtyqQuQyvbt21FYWCh1GV4THR2NmpoaqcvwKjrGwBfsxwfQMXZHQkICZs6c2WZ7yAaUzWbDggULpC7Da7KysrBixQqpy/AqOsbAF+zHB9Axdkd2dna72+kSHyGEEL9EAUUIIcQvUUARQgjxSyH7GVR7tFot5s2bh8TERLBsYGe3Xq/Ha6+95rXX53keZWVl+PTTT9HQ0OC19yGEhC4KqBbmzZuHUaNGQalUgmEYqcvpk/j4eFRUVHjt9QVBQExMDObNm4d3333Xa+9DCAldgX2a4GGJiYlBEU6+wDAMlEolEhMTpS6FEBKkKKBaYFmWwqkHGIYJ+EuhhBD/Rb9dCCGE+CX6DCqIrVu3Dtu2bQPLsrj//vtx2WWXtdmnsrISK1euhMViwaBBg7B48WIoFAoJqiWEkNboDMpPCYIAnud7/fyCggLs2bMHb7/9NpYtW4b3338fHMe12e+TTz7BLbfcgvfffx9arRbbtm3rS9mEEOIxFFB+pLKyEo888ghWr16NhQsXoqqqqtev9dNPP2HKlClQKBSIi4tDQkICzpw502ofQRCQk5ODyZMnAwCuvvpq7N+/v0/HQAghnkKX+DrwVW4dPj9SiyaX4LHXVMsZ3DM2CndkRHa4T3l5ORYtWoRhw4a1eezvf/87cnJy2myfOnUq7rjjjlbbTCYTkpKSxHFMTEybxRwtFgvCw8Mhk8k63IcQQqTiFwEVGxuLF198EdHR0RAEARs2bMC///3vNvs9/vjjmDx5Mmw2G1555RWcPn0aAHDjjTdi/vz5AICPP/4Y3377bZ9r+jq33qPhBABNLgFf59Z3GlAGg6HdcAKAe++916P1EEKIP/OLgOI4Dm+99RZOnz4NjUaDtWvX4tdff23VDmPSpElITk7GnXfeiYyMDGRlZeG+++6DXq/HggULsGDBAgiCgLVr12L37t2wWCx9qun2jAivnEHdnhHR6T5hYWEdPtaTMyiDwdDqEmF1dTWio6Nb7aPT6dDY2AiO4yCTydrdhxBCPE1gZOAVGvAKNQRW1uF+fhFQ1dXVqK6uBgBYrVYUFhbCYDC0Cqhp06bhu+++AwDk5uZCq9UiJiYGmZmZOHDgAMxmMwDgwIEDmDhxIrZu3dqnmu7IiOz0TEcKPTmDmjp1Kp577jnMnj0bNTU1KC8vx5AhQ1rtwzAMRo0ahb1792Lq1KnYsWMHxo8f7+myCSEhilPq4IhIhiMiGU5dAnhFOHiFBoJc1a3n+0VAtRQfH4+hQ4ciNze31XaDwYDKykpxbDKZYDAY2mw3Go0wGAxdvo9cLkdWVlarbXq9HvHx8X08gt7jeR5yudwjNSgUCvzud7/D4sWLIZPJ8PTTT6N///4A3JdKn3vuORgMBixZsgQvvPAC/vWvf2Ho0KGYO3culEplt9/HarW2+T76SlxcnGTv7SvBfozBfnxAaB2jxQkcr2VxzsLgnIVFjb1vCx/4VUCp1WosX74cq1evhtVq9ep7uVyuNg22XnvtNa+uX9cVlmXx5ptveqSG+Ph4zJw5s1WXyouv+/TTT4PjOFRUVEAmk7VaVLankyTMZrNkzdioEVzgC/bjA4L/GAUwuOWBLHy6Jx+2mKFAJ5fsmp/Eg3VawTqsYJ0NANLa3c1vAuriL8offvgBu3btavO4yWRCXFycODYYDDCZTDCZTMjMzBS3x8bG4tChQz6pmRBCQhWn0sMaPxbW+DH42ykFYBjRzk5OKM3noTSXQFlfAllTLVinFQxnR+tzq+vafQ+/CailS5eiqKgIX3zxRbuP7969G7///e+xdetWZGRkoLGxEdXV1di/fz8eeugh6HQ6AMD48ePx3nvv+bJ0QggJGa6wKFgGTofNkA60s3apor4EYVWnoKwvgaKhHIzQ+wUH/CKgRo8ejZkzZ+Ls2bP45JNPAADvv/+++FnM+vXrsXfvXkyePBlffvkl7HY7Xn31VQDuS0xr164Ve9pnZ2eLEyYIIYR4BqfQoCFlCqyJl7e5jBcuF4CCX6CuOAKFtdpj7+kXAXXs2DFMmjSpy/1WrlzZ7vZNmzZh06ZNni6LEEJCnsDK0Zg0AQ3Jk9vMvlPWnoOm7DCWzr8Zb/74o8ff2y8CihBCiH8RANhiM2BOuxa8St/qMUV9MfT5P0JpOQ8AkLM3e6UGCihCCCGt8HI16ofMhC02vdV2WWMV9AU/QlV9Br7onEcBFcS6027jzTffxNmzZyGXyzFkyBA8/PDDkMvlyMnJwfLlyxEbGwvAvZLHXXfd5etDIIT4mD0qDXXDbgav0onbWLsFuqKfoC4/AgaeXQKuMxRQfkoQBAiC0OuOtS3bbdTU1ODFF1/Eu+++Ky4Me9G0adPw+OOPA3CH1datW8V7p9LT0/H888/37UAIIQFBYOUwp10La/8rWm1Xlx+GPn8rWM7h85oooPxIZWUlXnrpJQwdOhT5+fl44YUXxDOYnuqo3cbw4cNb7Tdu3Djx6yFDhohLThFCQodTG4/aEbeB08SI21hHAyJOb0ZY9ZlOnuldFFAdiD75Ofrl/h0yl+dWtODkGlRl3Iua4fd0uI8v22205HK5sHPnzlbr/Z06dQqLFy9GdHQ05s+fj5SUlE6PjxASeJpiM1A37GaAbY4DVdVJRJz+FjKnd1f06QoFVAeiT/2fR8MJAGQuK6JP/V+nASVVu40PPvgA6enpyMjIAAAMGjQIf/vb36BWq/Hbb79h+fLldAM0IUFEANAw4Co0pE4VtzEuO/Rnt0BdedQnkyC6QgHVgZphf/TKGVTNsD92uo8v221c9MUXX6C+vh7PPPOMuE2j0Yhfjxs3Dh988AHMZjP0en17L0EICSACK0fdsFtazdKTN5oQdfzfkNtqJaysNQqoDtQMv6fTMx0peLrdBgBs3boVhw8fxssvv9xqQkZtbS0iIyPBMAxOnz4NQRDE5aQIIYGLU2pRmzEHTn2iuE1Vk4/IE1+D5ewSVtYWBVSQSktLw5VXXonHHnsMMpkMDzzwgDiD7+WXX8Zjjz2G6OhovPfeezAYDHj66acBNE8n37t3L77//nvIZDIolUo8+eSTYNpZd4sQEjic2njUjJzT6sZbTemv0Odv9en08e6igPIjcXFxeOuttzz2enfeeSfuvPPONttffPFF8euvv/663efOmjULs2bN8lgthBBpOfTJqBn1h+bligQe+jPfI7zcf7s/UEARQkiQs0ekonbUXRBk7makjMuGqBNfQVVbIHFlnaOAIoSQIGaPHIiakXMAmQKA+/6m6KP/gMJa1cUzpUcBRQghQcoWlYbakXPEe5xYuxkxR/8BeVPPOmdLhQKKEEKCkC1mCGrT72gOJ1u9O5z8aBp5VyigCCEkyNhihqE2/XaxsaCsqQ7Rx/4Bua1O4sp6hgKKEEKCiD0q7ZJwqkHM0X9AZg+8TuMUUEGsO+021qxZg9zcXHHliIULFyItLc3XpRJCPMCpjb9wWe9COFmr3eHksEhcWe9QQPkpX7XbAID58+dj8uTJfS2ZECIhV1gUakb9UbzP6eJnToEaTgDQu99+xCsqKyvxyCOPYPXq1Vi4cGGrtfR6qqN2G4SQ4MMpwlEz+m7wynAAAONsQnTOPwM6nAA6g+rQf8r/gy/Pfwkbb/PYa4axYbiz/524OeHmDveRot3GP/7xD/zrX//C6NGjMW/ePCgUiu4eEiFEYrxMiZpRfwCnjnJv4JyIPv4vKKyB39uNAqoDmyo2eTScAMDG27CpYlOnAeXrdht/+tOfEBUVBZfLhXfffRdff/01tXYnJEAIjAy1GXfCpUu4sIFHVN7XUJpLpS3MQyigOnBT/E1eOYO6Kf6mzvfxcbuNi9sUCgWuueYabNiwodP6CCH+QQBQN+xmOKIGitsiTn8raQdcT6OA6sDNCTd3eqYjBW+026ipqUF0dDQEQcD+/fupay4hAaIhZQpscSPFsbZgBzQVRySsyPMooIJUd9tt/O///i/q6+sBAAMHDsRDDz0kZdmEkG6wxQxBw8Dp4lhz/jdoi3+WriAvoYDyI1K023jllVc89n6EEO9zamJQN/xWcaysLYQ+f4tftGj3NJpmTgghAYKXh6E2Y454r5OsqQ5RJ74CI/ASV+YdFFCEEBIABDCoG3ErOE0MAIDhHIjK/TdYV5PElXkPBRQhhAQAy8CrYY8eLI4jTv4HikajhBV5HwUUIYT4uabYDDSmNC9Hpi3aA3VVnoQV+QYFFCGE+DGnxoC6oc33T6qqT0NbuFO6gnyIAooQQvwUzypQl3672K5dZq1CZN6GoJyx1x6aZh7EutNu49lnn0VTk/tD1vr6egwZMgTPPfcccnJysHz5csTGxgIAJk2aREsgEeJj5sG/gyvcAODipIh1YDm7xFX5jl8E1NKlSzF58mTU1tZi7ty5bR6/5557MGPGDACATCbDgAEDcOONN8JsNuPrr7+G1WoFx3HgOA4LFizwdfle4at2G8uXLxe/fv311zFhwgRxnJ6ejueff753B0AI6RNr7Eg0JYwVx/ozP0Bh7X2Hg0DkFwG1efNmfPnll61uIG3p888/x+effw4AmDJlCu666y6Yzc3dIR999FFxNYRAVllZiZdeeglDhw5Ffn4+XnjhBfEMpqc6arcxfPjwdve3Wq3IycnBwoUL+3IIhBAPcKmjYR56ozhWV+ZAXXlUwoqk4RcBdeTIEcTHx3dr3+uvvx5bt271ckUANmwAvvgXYPPgiuZhYcAf7gJmz+5wFynabQDA/v37MXr0aLGzLgCcOnUKixcvRnR0NObPn0/r9BHiAwIjQ2367RBkSgDurrj6M9+FzOdOLflFQHWXSqXCxIkTsWrVKnGbIAhYs2YNBEHAN9980+3VuOVyObKyslpt0+v1YlAaN20G78lwAgCbDeymzYh98MF2H+Z5HvHx8bjqqqvafXzp0qXdfiuWZRERESEej0ajQWRkZIf/ENi/fz9uueUW8XGdTocNGzZAo9Fg7969WLFiBdatW9fmeVartc330Vfi4uIke29fCfZjDPbjA3p+jOsKZKgwui/FyxkBi8fr0f/qxd4qzyO89XMMqICaMmUKjh071ury3kMPPQSTyYSoqCisWbMGRUVFOHKk6xV9XS4XVqxY0Wrba6+9hoqKCvfgplleOYPib5rV/B6XMJlMUCgUHT7ekzOomJgY5Ofni69VUlKCK6+8st3XNpvNyM3NxRNPPNHmcbPZjLS0NNjtdpw+fRp6vb7N45d+H30lKytLsvf2lWA/xmA/PqBnx9jUbzjqMn4vjjWnv8fnOw96qzSP6evPMTs7u93tARVQ7V3eM5lMAIDa2lrs2rUL6enp3QqoLs2e3emlOCl4o90GAOzduxfjxo2DUqkUt9XW1iIyMhIMw+D06dMQBAE6na7Px0AIaR+n0qN+WPP9TmGmPGjK/D+cvClgAio8PByXXXYZ/vu//1vcFhYWBpZlYbVaERYWhgkTJnSYxKGmu+02AGD37t1tzsD27t2L77//HjKZDEqlEk8++SQYJhSvghPifRebDwpyd8NSWVMdIk5tCsnPnVryi4B66aWXkJmZicjISGzYsAEfffQR5HJ3aevXrwcAXHXVVdi/fz9sLS65RUdH4/XXXwfgnn6+ZcsW/PLLL74/AA+Rot0GAPz1r39ts8+sWbMwa9Ysj9VCCOmYtf8VzZ1xBR6RJ78JqfudOuIXAbVs2bIu9/n222/x7bffttpWVlaGefPmeassQgjxOpc6BuaB14jj8JJ9UJpLJazIf9BSR4QQIhGBYVE3fLa4lJG8oQK6wl0SV+U/KKAIIUQiDSlXwqlPdA94FyJPbgza5oO9QQFFCCEScGrj0ZAyRRzrCncFfX+nnqKAIoQQHxNYufvSHuueWauoL0F4SeBO8PIWCihCCPExy4DprVYpjzy5EQwEiavyP34xi494Xn19PZ5//nmcPXsW11xzDR544IF297NYLFi5ciWMRiNiY2Px1FNPQavV+rhaQkKHQ5+MxqTmrgG6/G2Q22olrMh/0RmUnxIEATzf+w9LlUol7r77bsyfP7/T/b766iuMHj0a7733HkaPHo2vvvqq1+9JCOmcwMhQN2wWcOGmd1XNWWjKD0lclf+igPIjlZWVeOSRR7B69WosXLgQVVW97/2iVquRnp4OhULR6X6//vorrr76agDA1Vdfjf379/f6PQkhnWtInQpO0w8AwLjsiDj9bcivFtEZusTXgeJDDSj4pQGc03PXhWUKBgMnapGS2fElNE+12+iuuro6ccmjqKgo1NXV9ep1CCGdc4bHoiF5kjjWFWyHzG7u5BmEAqoDxYcaPRpOAMA5BRQfauw0oAwGQ7vhBPRssdjeYBiG1tsjxAsEMKgfelOLWXvFIb8QbHdQQHUgJTPcK2dQKZnhne4TFhbW4WPeOIOKjIxETU0NoqOjUVNTg4iIiF69DiGkY41J41vfkHtqM13a6wYKqA6kZHZ+KU4K3jiDGj9+PHbs2IE77rgDO3bswPjx4z3+HoSEsiqbe1r5Rdqi3ZA3VUtXUAChSRJB7P7778fatWuxfft23HvvvSgpKQEAvPPOOzh79iwA4Pbbb8eRI0fw8MMP4+jRo70+EyOEtCUA+HeBvMVae5XQluyTtqgAQmdQfsTT7TY+/PDDdrc/9thj4td6vR6vvPKKx96TENKsKX4MKswXzgME3t3jidba6zY6gyKEEC/gFOEwp10njsNL90PZUC5hRYGHAooQQrzAMug6CAo1AEDWVEttNHqBAqoFnuchCLQeVnf1dbULQoKVPSIVTXGjxHHEme/A8C4JKwpMFFAtlJWVweFwUEh1gyAIcDgcKCsrk7oUQvyKwLAwD7lBHI+N5qCqPSdhRYGLJkm08Omnn2LevHlITEwEywZ2dlutVpjN3rtLned5lJWV4dNPP/XaexASiBqTJjavVO6y49ZUBh9IXFOgooBqoaGhAe+++67UZXhEVlYWVqxYIXUZhIQUlyoCltSp4lhXuAsRV06XrqAAF9inCYQQ4kfMg2e0uOepAprzBySuKLBRQBFCiAfYYobA3q95Hc2IM99RE8I+ooAihJA+Elg5zIN+J47V5YehNJ+XsKLgQAFFCCF9ZEmZAk4dCQBgnFboz22XuKLgQAFFCCF94FJHo7FFnyf9uR/BupokrCh4UEARQkgvCbgwMULs81QCdcVRaYsKIhRQhBDSS/bowbBHD3YPBB4RZ76nPk8eRAFFCCG9IDAymAfNEMea8sNQNFZKWFHwoYAihJBeaEwaD04TDQBgnE3QFeyUtqAgRAFFCCE9xCm1aEiZIo51hbtoYoQXUEARQkgPWQZeA0GuAgDIG43QlB2UuKLgRAFFCCE94ND1R1P8aHGsP7uFVozwEgooQgjpJve08uYVI8JMeVDVFUpWT7Dzi9XMly5dismTJ6O2thZz585t8/hll12GFStWiL2Hdu3ahezsbADAxIkTsXjxYshkMmzcuBGfffaZT2snhISOpvgxcOoT3QPeBd25H6UtKMj5RUBt3rwZX375JV588cUO9zl69CiefPLJVttYlsWSJUuwaNEiGI1GZGdnY/fu3SgsLPRyxYSQUMPLVLAMvFoca0v2QW6rk7Ci4OcXl/iOHDnSq+Z66enpKC0tRVlZGVwuF7Zt24Zp06Z5oUJCSKhrSLkSvFILAGBtZoQX75W4ouDnF2dQ3TFy5Eh8+umnqKqqwttvv42CggIYDAYYjUZxH6PRiIyMjG69nlwuR1ZWlrfKlVxcXFxQHx9AxxgMAuX4qm3A8mMKXJwLcU+GBplXPd6t5wbKMfaFt44xIALq1KlTuO2229DU1IRJkybhjTfewJw5c/r0mi6XK6g7zoZCR106xsAXKMdXO+I2cLHuf/wqzKXYmv0xtnXzuYFyjH3R12O8OKfgUn5xia8rVqsVTU3um+D27dsHuVyOiIgImEwmxMbGivvFxsbCZDJJVSYhJAg59EmwxTZfmdHnb6P19nwkIAIqOjpa/Do9PR0Mw6C+vh55eXlITk5GQkIC5HI5rrvuOuzevVvCSgkhwUQAYB50nTgOM+ZCaS6VrqAQ4xeX+F566SVkZmYiMjISGzZswEcffQS53F3a+vXrcc011+C2224Dx3Gw2+3ibD+O47Bq1SqsXr0aLMti06ZNKCgokPJQCCFBxGZIh1Of5B7wLugKdkhbUIjxi4BatmxZp4+vW7cO69ata/exffv2Yd++fd4oixASwgRWDkvateI4vPRXmlbuYwFxiY8QQnytsf94cGERAADW0Qht8c8SVxR6KKAIIeQSnCIcDSlXimNt4S6wnF3CikITBRQhhFyiYcC0FquVm6ApPyxxRaGJAooQQlpwavrBmnCZONad20arlUuEAooQQlqwpF0LMO5fjcqafKhq8iWuKHRRQBFCyAX2yFTYY4a4B4IA/bkf6aZcCVFAEUII3DflWgY2TytXVx6DotHY8ROI11FAEUIILt6Ue6HXE+eErnCXtAURCihCCBEYWateT+Hnf4XM3vMWQMSzKKAIISHPmpgJTh0FAGCcVmip15NfoIAihIQ0XqaCJXWqONYV7aGbcv0EBRQhJKQ1pEyGoNAAAGRNddCUHZS4InIRBRQhJGRxKj0a+48Xx7qCHWAETsKKSEsUUISQkGUZcBUgUwAAFJYyhJlyJa6ItEQBRQgJSc5wA5riRotjHd2U63cooAghIcky8BqAcUeSqvoMVHVFEldELkUBRQgJOY6I5FZLGlGnXP9EAUUICSkCAPPAa8SxujKHljTyUxRQhJCQYo8ZAmdEsnvAc9AW/SRtQaRDFFCEkJAhgGm1pJGm7CDktjoJKyKdoYAihISMprhRcIXHAgAYlx3a4j0SV0Q6QwFFCAkJAiNDw4CrxHF46X7InFYJKyJdoYAihIQEa+Ll4MIiAACsoxHhpb9IXBHpCgUUISTo8TIlGlKuFMfa4j1gOYeEFZHuoIAihAS9xuRJ4JXhAC4uCHtI4opId1BAEUKCGqcIR2PSBHGsLdxFC8IGCAooQkhQa0i5EoJMCQCQNxihNh6XuCLSXRRQhJCg5VJFwJp4uTjWFe4AA0HCikhPUEARQoJWw4BpACsDACjqS6CqPiNxRaQnKKAIIUHJqemHprhR4lhXsIPaaQQYCihCSFBqGHAVwLh/xSlr8qGqL5a4ItJTFFCEkKDj0CbAZhghjvXUTiMgUUARQoKOZeB08eswUx4UDRXSFUN6TS51AQCwdOlSTJ48GbW1tZg7d26bx2fMmIE//elPYBgGVqsVK1aswNmzZwEAX3/9NaxWKziOA8dxWLBgga/LJ4T4EXtEKhzRg9wDgYe2cKek9ZDe84uA2rx5M7788ku8+OKL7T5eXl6ORx55BBaLBRMnTsQzzzyD++67T3z80UcfRX19va/KJYT4KQGtz57UFcegsFZLVg/pG78IqCNHjiA+Pr7Dx3NycsSvc3NzERsb64uyCCEBpnUzQhe0RbulLYj0iV8EVE/cfPPN2LdvnzgWBAFr1qyBIAj45ptvsGHDBgmrI4RIRQBgGTBdHGvKDkFupysrgYyZOHGiX9xWHR8fj5UrV7b7GdRFmZmZeOqpp/Dggw/CbDYDAAwGA0wmE6KiorBmzRq8+eabOHLkSJfvt2XLFmzbts1j9fubuLg4VFZWSl2GV9ExBj5PHt+hKhaf5bv/za1kBTw/1gmdwiMv3SfB/jME+n6Mw4cPb3f+QMCcQQ0aNAjPPvssnnjiCTGcAMBkMgEAamtrsWvXLqSnp3croFwuF1asWOG1eqWWlZUV1McH0DEGA08dn8CwMI17CNBEAwAUBT/jvR07+/y6nhDsP0Og78eYnZ3d7vaAmGYeFxeH119/HS+//DJKSkrE7WFhYdBoNOLXEyZMwLlz56QqkxAikaa40eAuhBPjbIK2ZF8XzyCBwC/OoF566SVkZmYiMjISGzZswEcffQS53F3a+vXrsWDBAuj1ejz55JMAIE4nj46Oxuuvvw4AkMlk2LJlC375hbpkEhJKBEYGS+pUcawt2QeWs0tYEfEUvwioZcuWdfr48uXLsXz58jbby8rKMG/ePG+VRQgJANbETPBiK/cGaM4fkLgi4ikBcYmPEELaw7OKS1q5/wyWd0pYEfEkCihCSMCy9r8CvFILAGBtZmrlHmQooAghAYmXqdCQPEkc64p+olbuQYYCihASkBqTJ0JQqAEAMmsN1JXHJK6IeBoFFCEk4HAKDRqTJohjXdEuMAIvYUXEGyigCCEBpzFlMgSZEgAgbzAizJgrcUXEGyigCCEBhVPq0Jg4ThzrCndSK/cgRQFFCAkoDalTANZ9C6fCfB6q6tMSV0S8hQKKEBIwXGGRsMaPFcd09hTcKKAIIQGjIXUawMoAAMq6IihrCySuiHhTj5c6uuKKKzB+/HiMHTsWcXFxiIyMhN1uR21tLc6cOYPffvsNe/bsEVcZJ4QQT3BqYtAUN1Ic6wro7CnYdSugVCoV5syZg1tvvRVxcXFgGPdfC4fDgdraWqhUKiQmJqJ///6YPn06Hn/8cezZswf//Oc/cfz4ca8eACEkNDSkXgUw7os+qpqzUJpLungGCXRdBtRNN92EBx54ADExMSgqKkJ2djaOHTuGEydOwGq1tto3NTUVGRkZmDBhAqZNm4Zp06Zhx44deOedd4K+YRchxHuc2jjYYtPFsa5gp3TFEJ/pMqCeffZZ/PTTT/j000+Rl5fX6b5FRUUoKirCt99+C41GgxmSMZcAACAASURBVBtvvBHz5s3DrFmzOmxIRQghXWnZyl1lOglFQ4V0xRCf6TKg/uu//gunT/d8GqfVasW6deuwceNGJCQk9Ko4Qghx6PvDHjPEPRAE6Ap3SloP6ZqC4ZGgdiJR40D/C/+ND3NCI+ehll38I0Atd6/+0dGMhS4Dqjfh1JLD4UBRUVGfXoMQErpanj2FGY9DYa2SrhjSRpiMR0zDKcwdWIWREU0Ypm9CbJgLbA9msPQ6oAghRCr2yAFwRA10DwQeusKfpC2IQMXyGB/TiAn9GpAR0YTBOhvk+csxbWjvXo8XOn6s2wEVFRWFMWPGICYmBk6nE0ajEQUFBTT5gRDiFQIAy8Dp4lhdcRRyW61k9YSyKKULUwwWTI21YHxMI1SyTlIFACcARpsC560KlDUpUWZVoNymhNnJosnFooljYePc/210sXjnD+2/TpcBJZPJ8MQTT+CWW24Rp5e3ZDQasXfvXmzatAknT57s3tESQkgX7NGD4dQnuQe8C7qi3dIWFGL0Cg4zEuoxI74eGZFNnV6yM6sSsTPfiuN1auTWq1HcqIJL6Ptdal0G1H333Ydbb70VlZWV2LVrF+rr66FUKnHFFVcgPT0dMTExuO2223Drrbdi3759WLlyJZ1VEUL6REDrz540ZYcgs5slqydUMBBwWbQVt/SvxfQ4S4dnSucaVNhj1OJgTThO1KvxyBPPYsXGFR6vp8uAuuGGG1BQUID77rsPNptN3O5yuTBixAjccMMNmDBhAm666SZMmjQJH3/8MZYsWYITJ054vFhCSGiw9RsBly7ePeCc0Bb/LG1BQS5K6cLN/etwc1ItkjTONo9zAnCsVoOfjDrsMelQalX6pK4uAyoqKgpbtmxpFU4tNTU1YefOndi5cycyMjLw8ssvY9WqVZg7dy6qq6s9XjAhJLgJYGAZcJU4Dj//K2TORgkrCl4JagfuGVCNm/rXtXu2dKI+DJtKI/FjpR5mp+/n1HW5WGx5eTni4+O79WK5ubl49NFHwbIsFixY0OfiCCGhpyluFLjwfgAAxmWDtmSfxBUFnzStDctGnce/p5zFHSm1rcLJ4mSxrjgK8/am4d5f0rC+NFqScAK6cQa1detWzJs3DxkZGcjN7bprZUVFBbZv347Jkyd7pEBCSOgQGNa9YvkF4aX7wbrav3pDem6Evgn/NciEqbENbR47UR+GL4uisaNSDzvvH40uugyozz//HNdeey3WrFmD//f//h82bNgAnuc7fU5TUxMiIyM9ViQhJDRY48eCU7t/dzBOK8JL90tcUXBIVDvw0BAjrk9oO9Hkt2oNPj3XDwdqwgE/Wx++y4Cy2+34y1/+gpUrV2LJkiX485//jK1btyI2Nrbd/dPS0vC73/0O58+f93ixhJDgJbByd7fcC7TFe8FyDgkrCnx6hQvz06rw+5QaKC45KdpVqcOnBf1wol4tTXHd0K0LizU1Nbj//vsxd+5c3H333bj77rshCO5rlp999hlMJhMcDgdiYmIwfPhwsCyLt99+26uFE0KCS2Pi5eBVegAAa7cgvOw3iSsKXEqWx50pNZiXVgW9ovUVr+0VOnx4NhaFjSqJquu+bn/yxXEcPvnkE/zf//0fpk+fjkmTJmHMmDFIS0tDWlqauF9hYSE+/vhjbN261SsFE0KCDy9TojG5+XNrbfEeMLxLwooC1/iYBjyVXt5munhOrRpvn45DTp1Gosp6rsdTMxwOB7Zs2YItW7YAcDczNBgMUCgUqK6uhtlMN9MRQnqmsf948MpwAIDMVgdN+WGJKwo80UoXFg2vwIxLPmcqblTi3dOx2GXUwd8+Y+pKn+cO2u12lJaWeqIWQkgI4uVhaEyeKI61hbvBCJ1PxCLNGAi4JakOjwytbHU5z+xk8eHZWKwviQLngWWHpECrmRNCJNWQPAmCPAwAILNWQV15TOKKAkea1oan08sxOqqp1fbvyyLw1qk41DoC+1d8YFdPCAloZgdg7X+FONYV/gQGna+UTQAWAv44oBoPDDFByTZ/v0qtCqw4kYAD1VoJq/McCihCiGS2lckgyGQAAHlDBcJMtIZnVxLUDrwwsgyXRVvFbU4e+EdBP3xyrp/f3GTrCX4TUEuXLsXkyZNRW1uLuXPntrvP448/jsmTJ8Nms+GVV14Ru/3eeOONmD9/PgDg448/xrfffuursgkhveRSRWCvsfmXqa5gZ4B9hO9rAmYl1uPxERUIlzd/1nSiPgyv5PQPiGnjPeXxqJ05cyZmzJiB8PDwHj1v8+bNePzxxzt8fNKkSUhOTsadd96J119/HVlZWQAAvV6PBQsW4L777sO9996LBQsWQKfT9ekYCCHe15A6VfzwXlFfAlXNWYkr8l9RShdeH1uK50eVieHk4oGPzvbDA/sHBmU4AV44g3r++echCAKsVivWr1+PL774ArW1XXfBPHLkSKeL0k6bNg3fffcdAPeitFqtFjExMcjMzMSBAwfE6e0HDhzAxIkT6T4sQvyYSx2DpvjR4lhXsIPOnjqQGd2Il0eXIkbFiduKG5V4Kae/X68C4QkeD6jvvvsODMNg0KBBuPvuuzFnzhxcffXVfX5dg8HQqhGiyWSCwWBos91oNMJgMPT5/Qgh3mMZMA1g3BdwlDX5UNUXS1yR/2EgYF5aFe4fbIKsRXp/VRyFd07HwcYFz2dNHelxQKWmpqKoqKjDx1999VXx6/DwcIwdO7Z3lXmZXC4XLxMGo7i4uKA+PoCOMVCVNjJYdVwhjh+dmoKUmcF1jC315meodFlwefHfEG8xidvsch1+S74f7JjRWOjpIvvIW39PexxQ//znP/HDDz/g5Zdf7nLfxsZG/PyzZzphmkwmxMXFiWODwQCTyQSTyYTMzExxe2xsLA4dOtTl67lcLqxY4fkWxf4iKysrqI8PoGMMVDUj7wJihgAARkfx+OLdNySuyLt6+jMcGWHFq2NKEaduXurpcI0GLx6LR5X9ewDfe6HKvunr39Ps7Ox2t/f4HNFisbS6pOYru3fvxsyZMwEAGRkZaGxsRHV1Nfbv34/x48dDp9NBp9Nh/Pjx2L+flugnxB859MmwXwgnCAJmJtF6e80E3JVajffGF7YKp0/PxeAvv6Wiyq7o5LnBqcdnUEeOHEFqaqrHC3nppZeQmZmJyMhIbNiwAR999BHkcnd569evx969ezF58mR8+eWXsNvt4qVEs9mMtWvXigmcnZ1N6wES4ocEAJaBzZ9HqytzEK8ZLl1BfkTJ8ngmoxwzE+vFbWYni5dz+uNnU+jOSu5xQH3yySd4//33MWLECOTl5XmskGXLlnW5z8qVK9vdvmnTJmzatMljtRBCPM8RlQZHZIp7wHPQFv0EgAKqn8qJ18eWICOyuXNwbl0Ynj+ahAqbUsLKpNfjS3zXXHMNDh48iLfeegs33nijN2oihAQZAYC5xdmTpvww5LY66QryE+kRTcieWNAqnDaWRuLhXweEfDgBvTiDutiskGEYPPfcc3j44Yexd+9e5Obm4uTJk8jPzwfHcV2/ECEkZNgMI+DSJbgHnBPa4j3SFuQHbkiowzMZ5VDJ3GvpuXhgzal4rCuOQqC1xfCWHgfUY489hmHDhol/kpOTcdNNN2HWrFkA3LPj8vPzkZeX1+ElOUJI6BAYFpYB08Vx+PlfIXM0SFeQxBgIeGSoEXMHVovbzE4WS48k4bea4Fjk1VN6HFCHDx/G4cPNzcRUKhWGDBnSKrQGDx6MYcOGUUARQtAUNwacJgYAwDiboC3ZJ3FF0lGxPJaNPo+r4yzitnMNKmQdSsb5JrqkdymPNCw8fvw4jh8/3vyicnmrNvCEkNAksHJYBkwVx9qSfWBdtk6eEbwiFS6suKwEo1r0btpj1GLZsf6wcjIJK/NfXlnN3OVyiSuNE0JCV2PiOPAqPQCAtVsQfv5XiSuSRpLGgTczi5Ec7hC3fVEYjbdPxYGnz5s61GVAqVQq2O32Pr2JJ16DEBJYeJkKDSlXimNt0W4wfOjdmBvdeBYfTihApNI9eYwXgDUn4/Dv4hiJK/N/XU4z/+qrrzBnzhwoFD2/i3nw4MF44403cPfdd/eqOEJI4GpMngRB4V5tW9ZUA03FEYkr8r2rYs2Ykv+GGE52jsFzR5IonLqpyzOo/fv3Y+HChbj33nuxbds2/Pjjj8jNze3wjCgxMRETJkzAzJkzkZ6eDqPRiM8//9zjhRNC/BenCEdj0nhxrCvcBUbgO3lG8LktuQZPjqjAxY7stQ4ZnjqUjNx6jbSFBZAuA+qVV17BunXr8OCDD2L27NmYPXs2eJ5HYWEhqqurYTaboVKpoNfrkZqaioiICDAMg9raWnzwwQf44osv4HQ6fXEshBA/0ZA6BYLMPStN3lCJMGOuxBX5koD5aVV4cEjzSuTFjUo8cTCFZur1ULcmSeTl5WHx4sVISkrCzTffjHHjxmHo0KEYNGhQq/3q6uqwa9cu7NixAzt27KAbdgkJQa6wSFgTmjsM6Aq2h8w0AAYCFg2vxF2pNeK2Gk0aHtguR73TK3PSglqPvmOlpaV47733ALgnPhgMBkRERMBut6O2thbV1dVdvAIhJNhZBlwFsO5p04r6Yqhq8iWuyDdkjIClI8taLfi6vyocldOzUO98S8LKAle31uL785//3GYFc7vdjtLSUuTm5uLs2bMUToQQOMPjYIsdKY7150KjlbuK5fHG2JJW4fRjhR5PHUoGJwuTsLLA1q2AeuCBB3Dttde22qZSqbxSECEkcFnSrgEYdySpqk5BaS6RuCLvC5dzWD2uCFfGNi/ftL4kCi8e7Q+nEPxt2b2p19+9e+65B5s3b273sZiYGKjV6l4XRQgJPPbIAbBHX/hcWuChK9ghbUE+oFdweGtcEca2WB1ibX4/rDgRTzfgekCf4j0iIqLd7bNnz8aWLVv68tKEkAAi4MLZ0wXqimNQWKukK8gHopQuvHNFIdIjmpduWnMyDn87GwtajdwzvDathGHoB0RIqLAZRsCpS3QPOCd0RT9JW5CX9VM58fa4IgzQupcu4gVgxYkEbCiNkriy4ELzHgkhfSIwbKtW7uHnD0BmN0tYkXfFhTnxzhWFSNK47+/kBODVnER8Xx4pcWXBhwKKENIn1oTLwKmjAVxsp7FX4oq8p7/agbevKEKC2h1OLh7472P98WNl+x93kL7pdkAJguDNOgghAYhnFWhIbdFOo/jnoG2nkayx4/9dUQRDmHvBWwfP4PkjSdht0klcWfDqdkDNnz8fU6dORV5eHvLy8hAfH+/NugghAaAxeSJ4pbsLLGurR/j5AxJX5B2XhpOdY/D04WTsr6YOuN7UrYA6cOBAq465s2fPFh979913cebMGZw9exZnzpxBfn5o3DVOSKjjFBo0Jk0Ux+4FYYNvebNLw6nJxeDJwyk4VBMucWXBr1sBtXjxYgBAQkICRowYgeHDh2PEiBEYOnQoxowZgzFjxoiXAHmeR1NTU2cvRwgJAg2pUyHI3TfsyxuNUFfmSFyR57UXTksOpeBwLYWTL/RokkR5eTnKy8uxfft2cVtSUlKb0NJqtfSZFSFBzKWOvmRB2B1gEFz/z1M4Sa/Ps/hKS0tRWlqKrVu3ittSU1MxYsSIvr40IcRPmQdeIy4Iq6wrgqr6jMQVeRaFk3/wyjTzoqIiFBUVeeOlCSESc+iTYTcMF8e6/G1BtW5CEoWT36CVDAkh3SYAMA9qXjg6rPI4lA3l0hXkYQlqB94ZR+HkLyigCCHdZjOMgFOf5B7wrqBaEDY2zIl3xhUhTu0OJxtH4SQ1CihCSLcIjAyWgc0Lwoaf/xVye30nzwgc/VTucEq8sHyRnWOQdSiZwkliFFCEkG5pTLwcnNq9GCrjtEJb9LPEFXlGlNKFt8cVITncvfCrkweeOZKEAzV0E67UKKAIIV3i5WGtljTSFe0Gy9klrMgz9AoX3mqxKrmLB5YeScYvVbR8kT+ggCKEdKkh5UoICncTUllTDTRlByWuqO+0cg5vjSvGYJ07aDkBWHaM1tbzJxRQhJBOucIi0dj/CnGsO7cdjMBLWFHfqWU83swsxjC9e2FbXgBeyUnE9kq9xJWRlvym3cbEiROxePFiyGQybNy4EZ999lmrxxctWoTMTPed62FhYYiKisKMGTMAAHv27BHXAKysrERWVpZviyckiFkGXg2w7l8VivoShFWdlLiivlGyPFZcVoxRLdq0v56bgB+on5Pf8YuAYlkWS5YswaJFi2A0GpGdnY3du3ejsLBQ3GfNmjXi17///e8xbNgwcWy32/HnP//ZlyUTEhIc+iTYYjPEsf7cjwF9U66MEfDqmFKMi7GK21blxeM/56kTrj/yi0t86enpKC0tRVlZGVwuF7Zt24Zp06Z1uP+MGTOwZcsWH1ZISOhx35R7vTgOM+ZCaS6VrqA+YiFg2ajzmBrbIG57/7QB64qjJayKdMYvAspgMMBoNIpjo9EIg8HQ7r7x8fFISEjAwYPNH9IqlUpkZ2fjww8/7DTYCCHd1xQ7Ek59f/eAd0F3bnvnT/BjDAQ8nVGO6xOaW9F/ei4GnxS0/3uG+Ae/uMTXE9dddx127NgBnm/+kPb222+HyWRCYmIi3nnnHeTn5+P8+fOdvo5cLg/qz6ri4uKC+vgAOkZvsnPA8qMKwH3fKq5LYjBr0oMefx+fHJ8gYFTZPzG4Kk/clB9zLXSj5yKL8f4FS/p72nt+EVAmkwmxsbHiODY2FiaTqd19r7/+eqxcubLN8wGgrKwMhw4dwtChQ7sMKJfLhRUrVvSxcv+VlZUV1McH0DF6kyV1KhoGXAUAYB0NyPn3u8jlHB5/H18c332DjBg8uEocbzofgdd+KIOA//Hq+15Ef0+7lp2d3e52v7jEl5eXh+TkZCQkJEAul+O6667D7t272+yXmpoKnU6HnJzmxmg6nQ4KhQIAEBERgdGjR6OgoMBntRMSbDilDg3Jk8WxrmAnWC+Eky/8IbUa97YIp+0VOryemwghoKd6hA6/OIPiOA6rVq3C6tWrwbIsNm3ahIKCAtx///3Iy8vDnj17ALgv77XsOwUAAwYMwNNPPw2e58GyLD777LNWs/8IIT1jGXg1IHP/o09uqYC64qjEFfXOTf1rsWh4pTjeZwrHsmNJ4AQKp0DhFwEFAPv27cO+fftabfvwww9bjf/+97+3eV5OTg7mzp3r1doICRUOXSKa4keLY33+1oDslDs9zoxnMprbgBytVePZI8lwUTgFFL+4xEcIkd6l08pVppNQ1Qde49ErYhrw0ujzkF3IotPmMDx5KAV2nn7dBRr6iRFCAAA2QzqcEcnuAc9Bf+5HaQvqhZERVrwxtgRK1n3WV9SoxOKDKWhwySSujPQGBRQhBAIrhzmtuVNu+PlfIbfVSlhRzw3S2rDq8mKo5e5wqmiSY9Fvqah1+M0nGaSHKKAIIWhIuRJ8WAQAgHU0Qlu0R+KKeqa/2oHV44qhV7jvj6yxy7Dot1RU2hQSV0b6ggKKkBDnCotEQ/Ikcawr2BFQvZ5ilE6sHleEfip3q/YGJ4vHD6ag2KqSuDLSVxRQhIQ486DrmlcrN5dBXXFE4oq6TyfnsHpcMZJatGp/6nAyTlvUEldGPIECipAQZo9Kg73fcHGsP/t9wNzCGibjsTKzueGgiweeP5qEI7XhEldGPIUCipAQJTAs6gfPEMfqiqNQWsokrKj75IyA18aUYHSLnk5/PZ6IPdQNN6hQQBESohr7XwFO0w8AwLhsAbNaOQMBL4w6j0mGRnHb/+bF4XtqOBh0KKAICUGcUouG1ObWNNqi3ZA5Gzt5hr8Q8MSICsxo0TYjO78f/l0cI2FNxFsooAgJQZaB10CQu2e5yRtNCD9/QOKKuufeQSb8PqX5/qx1xVH48Cz1dApWFFCEhBiHvv8l6+1tASPwnTzDP/w+pQb3tViZfEu5Hm/mxQMBM62D9BQFFCEhRACD+sE3iGOV6SRUtf7fnmZGQj2WjKgQx/tM4Xglpz+1zQhyFFCEhBBr4uVw6RLcA84Jff7Wzp/gByb1s+CFkc0NSHNq1XjuKK1MHgoooAgJEZxSC8vA6eJYW/wz5PZ66QrqhtGRVrw2thTyC7+p8i0qPHk4BTaOfnWFAvopExIizGnXQZCHAQBk1mpoS/Z18QxpDdLa8D+ZxQiTuRd/LbMqsPhgCsxOWpk8VFBAERIC7JEDYYsbKY4jznwHRuAkrKhz7S7+ejAVVXZa/DWUUEAREuQERob6Ic0TI8Iqj0NVVyhdQV1ob/HXxQdTUWpVSlwZ8TUKKEKCXEPyJHAa942sjMsG/bltElfUsY4Wfz1jCZO4MiIFCihCgpgrLBINKVeKY13BTsgcDRJW1DEVS4u/ktYooAgJUgIA8+DfATL35zZySzk0ZQelLaoDMkbAa2NLafFX0goFFCFBytZvGOwxQ9wDQXBPjIAgbVHtYCDghZHnMdnQfGZHi78SgAKKkKDEy1Tus6cLNOWH/LOVhuBe/PV3ic2Lv66lxV/JBXKpCyCEeJ5l4NXgVXoAAOtogK5gh8QVtW9E5XoMb7H469fFUfgbLf5KLqCAIiTIOPTJsPYfJ471Z7eAddkkrKh9f0itxvDKE+J4a7keq2jxV9ICXeIjJIgIjAx1Q2eJY1X1aYSZTnTyDGncmFiHRcMrxfFekxYv5/QHT+FEWqCAIiSINKRcCS78YpdcOyLOfO93v/KnxZrx3Mjmz8OO1qrx3JEkWvyVtEGX+AgJEk6N4ZJ7nrZDZjd38gzfuzy6Ea+MOQ/ZhSyqC0vGk4c0sPP0b2XSFv2tICQICGBQP2wWwLoXUlXUl/jdPU8j9E1447ISKFn3VPeSRiX2pj2JBhct/kraRwFFSBCwJl4Opz7JPeA5RJze7FeX9tK0NqweV4RwuXvxV6NNjoW/pcCuiJC4MuLPKKAICXCcSg/LwKvFsbZ4DxTWqk6e4VtJGgfWtFiZvM4hw6LfUlFho8VfSecooAgJYAKA+iEzIchVAAB5owna4r3SFtWCQeXEWy1WJm90sVh8MAWFjSqJKyOBgAKKkADWFDem9XJGpzf7TZ+nKKULb48rQoLavTK5jWOw5GAyTpnVEldGAgUFFCEBilPpYR58vTjWnD8ApblUwoqaaeUcVl9ejFStAwDg5IFnjyTjaB2tTE66z2+mmU+cOBGLFy+GTCbDxo0b8dlnn7V6/MYbb8Rjjz0Gk8kEAFi3bh3+85//iI/Nnz8fAPDxxx/j22+/9WnthPiaAKBu6E0tWrjXQF+wXdqiLgiT8ViVWYyhevfqFZwALDuWhF+qtBJXRgKNXwQUy7JYsmQJFi1aBKPRiOzsbOzevRuFhYWt9vvxxx+xatWqVtv0ej0WLFiABQsWQBAErF27Frt374bFYvHhERDiW00Jl8ERneYeCAIiT20Ew7ukLQqAkuXxxtiSVm0zlh9PxI5KvYRVEW8TBAGCQoCgufBHeeGPSgCUcH+tEAAZ3NftWEBgBffXnSyw7xcBlZ6ejtLSUpSVue8u37ZtG6ZNm9YmoNozYcIEHDhwAGaz+4bEAwcOYOLEidi6das3SyZEMi5VBMxp14nj8NJf/OLSnpwR8NcxpRjfr1Hc9r95cdhcRm0zAp3ACBDCBfB6HkLEhf/qBPAaHkK4gFctr8J5r9Pj7+sXAWUwGGA0GsWx0WhERkZGm/2mT5+OsWPHori4GGvWrIHRaITBYEBlZWWr5xoMXa+GLJfLkZWV5ZkD8ENxcXFBfXxAaB4jLwDvnZTDZHZ/fBwbJmDJ7ZdDyV4uVYkAAEbgMK7ofSTVN/d0OhF/OwaMuQWd/YRC8WfozziBQzVfDSNvRCVXCSNvhJE3oo6vA4eOJ9844flwAvwkoLpjz5492Lp1K5xOJ2699Va88MIL+Mtf/tLr13O5XFixYoUHK/QvWVlZQX18QGgeY2PiOJiH3OAeCDy4vR9j9Q/S9nlyNxwsQ1L/enHbJ+di8P4PJwGc7PS5ofgz9BcCK4CP5sHH8uBiOfD9ePBRvPsyXG+4AMbKuP/YGTAO93/hgPtrBwNwAHj3H4Zn3F87AYxv/yX9IqBMJhNiY2PFcWxsrDgZ4qKLl/AAYOPGjXj00UfF52ZmZrZ67qFDh7xcMSG+5wqLgiXtGnEcXrLPD5oQCngqvQIzW4TTv4qi8f6Z2E6eQ6TAa3hwCRz4+AuBFMP3KAEYKwPGzIA1s83/bWDAWlksun8R3vqft8B4eP0SvwiovLw8JCcnIyEhASaTCddddx2WLVvWap+YmBhUV1cDAKZOnSp+PrV//3489NBD0Ol0AIDx48fjvffe82n9hHibAAZ1w2+BIHOvviBvNEJX+JPkVS0cVonbkpsbDn5TEonVJ+NAPZ2kx2vdgcQlcuASOAiRncxGaIGxMGBrWbDVLNiaC3/MLBhnxz9TNaP2eDgBfhJQHMdh1apVWL16NViWxaZNm1BQUID7778feXl52LNnD+bMmYMpU6aA4ziYzWa8+uqrANxnVmvXrkV2djYAIDs7u9XZFiHBoCHlSjgjkt0DgUfkyY0S35Ar4MHBJvxxQI245buyCPzPiQRQOElDUAjuMErm4Ep2QYjoOpAYMwOZUQbWyEJmkoGtZt2X4vyEXwQUAOzbtw/79u1rte3DDz8Uv37vvfc6PDPatGkTNm3a5NX6CJGKQ5eIhgHTxLG2aDcUDRUSVgTcO8iE+YOa1/vbXqHDX48nUsNBHxIggI/hwaW4A4mP6+LzIxfcQVQmg6xSBplJBsbm3z8vvwkoQkhbNg6oG3EbwLhn7SnqS6At2iNpTfPTTLhvcHM4/WzUYtmxJHDUcNDrBFYAF8+BG8jBNcAFQdfJWZITkFXIICuXQVbmPkti+MD6GVFAEeLH1hfKwKmjAACMy4bIvG/AdHZno5f9aWAVHhzSPIFpnykcz1I3XK8SZIL7LGmgC65UF9DJOrtsFQtZicz95B9grwAAIABJREFUp0IWcIF0KQooQvxUk2EEfq1qvmYTceY7yO31nTzDu/6YWo1Hhjbfr7i/KhzPHkmGU6AlPT1NkAnuz5IGXQiljjqT2AF5iRyyYncosU3B9bOggCLED3EqPeqH3CiOwypzoDbmSlbPnJRqLBzefEP8b9UaPH04mVq1e5DACuCSOLgGu+Aa0HEoMRYG8iI5ZIXuS3eBfpbUGQooQvyMe0r5bAgKd1sKWVMdIs58L1k9dyTX4PERzeF0uEaDpw6nUDh5gAABfD8ermEuuAa7IKjbv3zL1DGQ58shL5CDrWK9MqXbH1FAEeJnGpMnwRGZCsC9SkPkyW/AcnZJarkzpRpPtAino7VqPHkoGTaOwqkv+HAerqEuOIc4IUR3EEr1F0IpX+6e/h0iodQSBRQhfsShT4Jl4HRxfH1/Hjm7pFkI9g+p1VjU4rJeTp0aTxxMgZXr7Vo4oU1gBXADODiHO8Elc+3eLsY0MpCfkUN+NrTOlDpCAUWIn+AUGtSm395qSvmM8XHIkaCWuwdU4S/DmidEHKtV43EKp14xckbYJ9nhHOoE2msm7ATk5+SQn5FDdl4GhmZEiiigCPEDAoC64bPBq9x9kxinFVF56yGb8ZDPa5k7sAqPtpitd6RWjSUUTj0iyAS40lxwZjjxTuM7wJi2+8hKZZCfcn+uxLgolNpDAUWIH2hImQJH9CBxHHlyA2R23y/ZNW9gFR5uEU6HazRYcigFTfSZU7fweh7OdCecw9o/W2IaGMhPyqE4pQBroe9pVyigCJGYPXJAq6WMwot/RlhNvo+rEPBfaVV4oMVNuAdrNHjyUApNiOiCwAjgUjk4My58tnQJGWRAPqA4qYCslC7h9QT9zSNEQpxSi7oRt4qfOynriqAr2OnjKgQ8PMTYKpx+q9ZgyUEKp84ISgGOMQ5Y/2iF7QZbm3BiLAyU+5VYol0C9VY15CVyCqceojMoQiQigEHdiNvAK7UAANbRgMi89T5dyoiBgMeHV+DO1OaWGfurwukm3E5wURyco5xwDXEBikseFABZsQyKXAVkJe6zJe31WknqDAYUUIRIxDLgKvF+JwgCIvO+gczR0PmTPIiFgGcyynFzUp247SejFi8cTYKDwqkVAe718JyjneCS2mlz0gQo8hRQ5NFnS55EAUWIBJr6DUNj6hRxrC36Caq6Qp+9v4wR8OKo85iR0DwRY2u5Hi/l9KdVyVsQ5AJcQ11wjHa02/CPrWKhyFFAflYOhqPvm6dRQBHiY85wA+qHzxbHypp8n7bQUDA8XhlzHlfFWcRtm0ojsTw3gfo5XcBreDgznHBmOIGwSx8EZAUyKHOUYCvoZlpvooAixId4uRq1GXPE1u2yplpE+fBzJ42Mw/KxpRjfr1Hctq44Cm/mxUOgX7Tgo3g4xjjcny9detuX/cJlvOMKsA10Gc8XKKAI8REBDGrTb2/R38mOqOP/Auuy+eT9IxQuvHl5MdIjmt/vs4IYvHs6FqHcpl2AAD7BHUzcgLafLzFmBoocBRQnFWCcoft9kgIFFCE+Yhl0HRxRA8Vx5MkNUFirOnmG58SFObH68iIM0DrEbR+cMeDjc/0QquEkMO618RxjHe526ZdgK1gojyohK6R7l6RCAUWID1jjRqMxaYI41hbuQlj1aZ+8d2q4HWsuL0Kc2gUA4AVg5Yl4rC+N9sn7+xuBvTDxYWw7Ex8EQFYocwdTBS3tJDUKKEK8zKFLRP3Q5uaDqqqT0Bbt9sl7p0c0YVVmMSKV7ktXTh7472NJ2F6p98n7+xNBIbiXIRrthBB+STC5APlpOZRHlWDr6fMlf0EBRYgXcSo9ajPuBFj3/2ryRiMiT270yUW1K2Ia8PrYEmjk7l/GVheDZw4n40BNaN04yofxcI5ywjnSCaguedAOKHIVUOQogq5dejCggCLES3iZCjWj/gBe9f/bu/f4qOo7/+Ovc5nJPZM7CRBA0AgBQW2xaCtKNz+7paLV1m6rVNtf131sq4+Ki6W7q2C7W9utlmrd3dauNl1K258/adcWqJdfsa4GTdUV5RpBICAhCRlym0wmmTlzzvf3x0wmGRIgQCaZGT5PH/OYmXPOzPkeh8x7vud8L3kAaFYfhbs2oNuh07zy3C2d3MU/zG3GjH7ndoUMVm6bxp7ukeZ7SE9OjoN1qYU12xo24oPWq+Ha4cK1Rxo+JDMJKCESQGk6nXM/SzinLLLAsSnc8xvM/s5Tv/Dc98z/nnWcOy8cHFfvWJ/JPW9P53DvidWH9OR4HEKXjdxUXOvScL/rxtxnojkSTMlOAkqIMaaA7qrr41vs7d1ERtfhhO7X0BSrqlu4YcjQRe/3ZLDy7Wl4gycOGpd+7GIb63KL8MzwsIaJulfH/Y4bo1Fa5KUSCSghxph/xjX0lc+PPc9tfJmstl0J3We2YfOdBU1cWTrYAffN4zn84/ap9IbTuzWaPckmdHkIe/rwPkx6czSYjhgy4kMKkoASYgwFyhfgn3517HlWyzvkfvBaQvdZ7LZY+6EjXJw/2AH3uaMevrd7MuE0rS0oFPaUSI3JnjLCHEyHjUgwSVPxlCYBJcQYCRbOpLvqU7HnGR0H8Lz/fEJ/t1fl9fHwZUdifZwAag+U8OT+UtKxA+7AqOKhD43QuVaBcdDAvc2N0S7BlA4koIQYA6G8yXRWfyY28aDZ00rBnt+iqeEjFIyVa8p8PHjJUbKizcjDDjy8p4JNRwsTts+JojRFeGYY6zILp+SE/6cOmO+buN9xo3dJU/F0IgElxDmycsrouOQLKDPSSk7v76Zo19MJbE6uuP2Cdr5a1RZb4rd0Htg+lTfa06uPk9IV4QvDhC4LoQpP6Fxrg/meiftdt8zBlKYkoIQ4B+GsIjrm34pyRfoX6aFeinb+OmETD7p1h7+f28InJ3fHljUFXNy3bVpaNSNXhiJ8cXQ4ovwTgskC1x4Xru0u9IAEUzqTgBLiLIUzPLTPXx6bsl0L91O089e4Au0J2V+hO8z3Lz3CJYV9sWXbOrL5h3en4rPS409ZmdHhiBaMMBxREFy7XLh3utH60+/6mhguPf5VCzHObHcuHQtuw8mMjGmn2SGKdj6Ny38sIfur9vTx3QXxjSF+d6SAtQ0VadFST7kV1jyL0CUhOHGwiz5w73Tj2uVCC6X+sYrRS5qAWrRoEStWrMAwDDZu3Mj69evj1n/+85/nhhtuwLZturq6eOihh2htbQVg69atHDhwAIBjx46xatWqcS+/OH84ZhYd82/FzoqOBu6EKdy1AbevKSH7Wzalk/uqW3HrkRqFreBf907i/x4uItVb6vkdP8ErgpGZa084Q6n1ari2R4cjCqf2cYqzkxQBpes6K1eu5J577qGtrY3a2lrq6uo4dOhQbJt9+/bx5S9/mWAwyE033cRdd93F6tWrAQgGg9xxxx0TVHpxPnHMTDrmf2FwCCPlULjnv8joahzzfbk0h0uP/Jyb5rXElvksnQd3TOXPx1O7McTAOHmP+h/FutyKW6f1aLjfcWPuNdFsCabzWVIEVHV1NU1NTTQ3NwOwZcsWFi9eHBdQ27Ztiz3evXs3f/mXfznexRTnOduVTcf8WwnnlkcWKEXBe79PyLxOpRkW37v0CBd0vBdb9n5PBn//TiXNfe4x3994cfKj4+RVjTBOXmc0mPbLOHkiIikCqrS0lLa2wSazbW1tzJ0796TbL1u2jPr6+thzt9tNbW0ttm2zfv16Xn311dPu0zTNtD4VOGnSpLQ+PhjfY+wOwU8aXISjF+c1FJ+babNo0aeAT536xWeoxN/AFYd/TMaQqeCPFCyiYd6XWf6x1Gyp12K3UBesY3d4N4r4xg8VegWLMxYzJ28O+vT0a5Unf4tnLykC6kx84hOfYPbs2Xzta1+LLbv55pvxer1MnjyZf/u3f+PAgQMcPXr0lO8TDod5+OGHE13cCbNq1aq0Pj4Yv2O0M/Jpn78cOzt6zUk55O/dxKuv7OT0P4VGT0fxpVnHuWGWFyNagXDQ+VFDKc980A38aAz3lngKhVMRqTHZ00YYJ69F59aZt/Lso8/yh+h/6Uj+Fk+vtrZ2xOVJEVBer5eysrLY87KyMrxe77DtFi5cyJe+9CW+9rWvYVlW3OsBmpub2bZtG1VVVacNKCFGI5xZQMf85dhZBZEFyqGg4XdkefeM6X6K3Rbfmn+UDxcHYss6gga7qr/BMy9uHNN9JVpsOKLLQzjlw0fSMI5EhyNqMahaVSWDuIqTSor6dENDA5WVlVRUVGCaJjU1NdTVxU+JXVVVxapVq/jGN75BZ+fgnDp5eXm4XJGpBDweD/Pnz6excewvWIvzTziriPYFtw+Gk2NTuPu3Yx5OVxT7+cVVB+PC6Z2ObL5UP5P23Nljuq9EUrrCqrLo+1wf/Uv748NJgXnAJOs3WWT9IQujRcbKE6eXFDUo27ZZu3Ytjz32GLqus3nzZhobG7nzzjtpaGhg69at3H333WRnZ/PQQw8Bg83JZ8yYwTe/+U0cx0HXddavXx/XuEKIs2Hllkdmw412wsUJU7h7A5kdB8ZsH4amuPPCNu6YOdix11Hw8wMl/PxgKXaK9G9SLoU1x8Kab6FyRxiOaF90OKLupPg9LFJIUgQUQH19fVzDB4Ann3wy9vjrX//6iK/buXMny5cvT2jZxPmlv2gWXdWfQRnR1nK2RdGuZ8a0KfnU7BBrLjnKJQWDo0IcD5p8a8cU3u7IGbP9JJKT5WDNs7DmDe/DRAhcDS5cO1zovRJM4uwkTUAJkQwC5ZfRXfXJ2KjkmtVH0a5ncPuOjNEeFJ+e2sXXL26NjUIO8MbxHL69cwqdoeT/k3QKHEILolOqn1BcLaDh2unCtVtGfRDnLvn/GoQYBwrwz7gW//SPxZYZ/V0U7vw/Yza2XpE7zD/Oa+ajpYMDyYYdeHJ/Gesbi1FJ3Fgg1iJvQQh7xvAWeVq3hvtdN+Y+6Vwrxo4ElDjvKU2nu+r6uGnazZ4WinY+jWH1nuKVo3dNmY+/n9tCgXvwy73R7+ZbO6awr+fEweeSh9Kj8zDNt3DKhrfI04/puLe7MRoNtBS5ZiZShwSUOK85Zhad1TcTKrwgtiyjfT8Fe36L7lineOXo5LvCrJh9LG56DICnDxXxxPtlBJ3kvD6jMqKjis8bYVRxBcYhA/d2N3qrLs3ERcJIQInzlpU7ic65t2BnFsSWZbW8g2ffc2gnjHZw5hQ15T7und1KUcZgrelYn8l3dk3mfzqScyw9p8AhdEl0KCLXCSvDYO41ce+QFnlifEhAifNSoOwSuquWgjH4LZzb+N/kfrD1nOsDZZkW35jTwsfK4ictfKHZw9qGcvzh5OoDpDSFXWljXWJhV45wfSmg4doVHVVc5mES40gCSpxXlKbjm1lDYOoVsWVauD866Ov75/TeGopPV3ZyV1UbOebg9ZpjfSaPNFTwmjfvnN5/rCm3wpptYc21UJ7hNUb9uI5rh0sGbxUTRgJKnDdsVw5d1TcTKpgeW2b2eincvQGzr+Oc3rsqr4+Vc1qZP2S2W4DfflDIj/eVEbCTp9ZkF9lY86xIM/ETT+MNXF/a4UZvketLYmJJQInzQrDgArpm34CTMViLyfQ24Nm7Cd0OnfX75rts/ubCNj5d2Rkb4BXgsN/Nd3dPZkdX9rkUe8woQxGeFcaqtkYcH49+cL0X6b+k98j1JZEcJKBEWlOaQc8F19JbeeWQhYq8xpfJOfL6WdcPdBTLpnbxtxe1xTUdtxz4ZWMJ/3mwhFAStNBzPE6kNd7FFmQOX68f13Htip7Gk1lrRZKRgBJpy8ouoWvOpwcnGAT0UC8F7/2ejM6DZ/2+8zwB7p3TSrWnP275n4/n8MOGco4EJnbOJmUowheECc8JY08Z3ugBG8yDZqS2JM3ERRKTgBJpRwGByR/CN7MmrpVeRvt+PHs3nXXn28rsIH97URsfL++JW94ccPHYe+XUeXNhAr/s7SKb8Jww1kUj15Y0n4ZrjwvzPRO9f+Jrd0KcjgSUSCt2Rj7dFy0lWHzh4EInTP6BLWQ3/89ZxUehO8xXZnm5cWon5pDv9aCt8YvGEn7VWDxhHW6VWxG+MIw1e+SRHnDA+MDAtduFccSQ2pJIKRJQIi0oNAJTFtJzwbWDo5ADpv8YBQ3P4gocP+P3zDIcvjCjndtmHCfbjG+G/ceWfH7yfhktfe6TvDpxlK6wp9pYF1uRcfFGaCCo+TRc77kw95oymrg4c0phAhlK4QbcSuFW4ELhii5zKYUBuBR43nqLJf1BTCLLDBWZbNBARe6jzzUi3TEGfiZpQPgUP5okoETKs3In0V31Kay8yXHLc5reIO/gn9DUCNdhTiHbsPnMtE6+MKOdQnf8a9/uyObf906iwTe+4+cpFE6xQ7gqTPiiMCp7hJEuoteWzAYTo1lqS+elaHjkKEW2UmQ7kfusgefRx5lKkcmQx4rovcKtFBmM+Lvn5P7rWb6YgMORgBIpK2iDb+Zf0Dv1I7HpMQDM3jY8+57D7Ws6o/fLNmxumdbBF2Z04DkhmA70ZPDv+8qoPz6+15na7XZCHwphXWihCkcefklv0zH3mrj2u9CCEkppQykygTzHIU8p8hxFnnLIiT7OVYpc5ZCrFDmOIkdFbid2bUtlElAi5Sigr2we/7LDFd983AmTe7iO3CP1aGqE6zEnkWPafG5aB5+f0U6+K/51LX0ufra/lOebPTjjFExOjkN4ZqSm9KPeH8HC4dtovRrmPhPXPhd6p5zCSxnRWoxHKTyOg0cp8h2HfKXIdyKPPdHgyXcitaGJEgaCmkYICGkaIQ0sNCwNQtH7MBphDarnL2Dbzh2E0bABWwMHsNFwACf6XMVukb8lB7A0+PJJyiABJVJKsGA6vpk1hPMqYEj/WndnI573n8Ps6xz1e5Vnhvjc9A5umNoVNzQRwNGAi3UHS3i+uYDwOEwj4eQ4hGeFCc8Mj9yRFiAE5iETc5+JcVSmt0gq0VNnhUpR4DgUONF75TDll7/iH309sVBKdOhYQK+mEdA0ejWNvujjPk0joEfu+9Dojy7r16BfizwPahpBIsFka6P/97XqMzfzywP7z7rMElAipVnZxfTM/AuCxVVxy/VQL3kHt5B1bOeo6zfVngBfmN7BknJf3OgPAE0BF/95oJQXWjzYCQ4AJ8+J9FeaFcaZNHIoGRjQCOZ+E/OwdKadCJpS5ClFoeNQ5DgUOopC5VA48NiJPD5p77fduznbURhDgE/X6NF0/JpGT/Rxb/SxX4vcejWdXj0SSCGAMwiXZCYBJZJaOLOA3sqrCFRcGnedCdvif1Xq7Hjm30c1VJGpKRaX9fC56e0sOGG8PIBDfjfrG0t4MYHBpFA4pQ7hGWHsGTZO8UlqSg4YzQbmAZO/u/7v+NcX/zUh5RFA9LpNkeNQ7ETuB26FyqEoWhMay+s6/UC3rtOta3RrOj5dwxe9H3jeo2n4dJ0gpE3YnA0JKJGUwlnF+Kd9lL5J8+KDSSmyju0gr/G/WXrVV9l1mnCakhXihqmdXD+lK25epgFvtefw60NFvHE8NyFTritTYU+2CU8PY0+3UbknmWfKAeNoJJTMQ2ZsWousZck7224qMGLhEwmb4mj4DL0fq3E/gkCnrtOla3RpOl3Rx1ffeCPrN/8hEki6Tv95HDhnSgJKJBUrZxL+6R+lv2TOsF+O7s5G8g9swdV77JTvMVBbunFqJ1eUDB81wnLg/7V4ePpwMft7Rhhy4Rw5HofwtDB2pY092T75X1k4GkoHo6EkLfDOWJZSFNuRoBl6K3IcSqKND8aiCYlf0+jUNTo0nU594KYNPo5ezxmptnPpggXse/HFMSjF+UcCSkw4hUawaBaBKR8mWHThsPXuzkZyP3gNd9ehk9ZxNBSXFPTxiYpuPl7uixvAdUBbv8mmpgJ+11TI8eDYnbRRGdFa0tQw9lR7xLmVYvrBPBwJJOOIIdeUTkFTinylhoVP7GY7jMVY8f1Ah67Trut0RG+duhZ73KHrhKTWMyEkoMSEsV3Z9JUvIFDxIeysgmHrM9rfJ/eDrbh9R0/6HjNz+7muopvrKnxUZFnD96Gg3pvL75sKqT+eOybXl5SpsCfZ2FNt7Ck2Tqlzyq5ReoeO8YGBediMDM4qre+AyOk3V3s7sy0rGjpqWC3oXH9GOECXpsXCp31IEEUen7zmIyaeBJQYVwqw8qcSqLicvrJq0E/4J6gUmd4Gcj947SSn8hSz8/tZXNbDX+y9n5s+OnJ4tfaZbD5awKajhbT1n9vXnHIr7HIbuyJyc0qdU3ezt8BoMjA/iNSSdP952E8p2t+n+ITrPsNOv/1gLavOYTchiIXO0PA5Hg2fLl0/o+bSIrlIQIlxYWWX0jdpHv2lc0esLWlWgOzW7WQ3b8Psj+/LZGqKy4p6WVzWw9WlPUzKCkdWxM92gc/S+VNrPi+2eNjemX1WjR4UCpUXCSSnzIncFzuc8kKGA7pXxzhqYDQZGK1G2k+R7hrS7Hpo44OiIWE0Flf3/NHaT7uuDQuidl2nR2o/aU0CSiRMOLOA/tJq+srmEs6dNOI2Lt9RspvfJsu7B82JBg+K6TkhFhb38pFiP5cVBYZ1pB0QtDXq2vJ4scXDn4/nnnGnWpWhsEuiYTTJxpnkoLJOcQ0pSuvQMI9GOswazQZaKH2+JE2lYn18BptbD4ZRUXTonXPlALbHwyG/PxI4xvCakLR4O79JQIkxo9AIeSoJFl9IsOgiwjmlI26nWX1kehvIadmGy98KKMozLeYX9vHhol4WFvspH6gljcBn6Wxty6Poo7fzD0/9kX57dKfQnEwHp8TBKXUioVTqoPJH8UXrgN6uY7REwshoNWLNwFNKtM9PQXS0g6LoiAeF0RAqGMPwgUiz6+HXfHQ6orWhTl1n5Te/ycMPPzwm+xPpRwJKnDUF2BkeQgXTCRbNIlg4E+U6Sb8d2yKzfR9ZbbvI6jjAhbkBFhQFuGRmgAUFgcHTdifRHHBR583j1bY8tndmYyuNVUs/RL/90vByGQqn0MEpdnCKBu9HHAF8JEEwjhnobTrGMQPjWHLXkDQVGTjU4yg8yqHAcfA4igIVvXecWCiNVdtFGyLNrLX41m8DAdShR0Y7kNNv4lxIQIlRU4CdVUyoYBpBzzRCnmk4mZ6Tv8C2yO46wBzfn5nvvMecXD+zZ/Vx0aX9ZJmnDgu/pfN2Rw5vtkduTQE3Q5vKKRQ9Tk+k0UKBE7kVRu5Vnhr9gON2pHakHx8SRl3ahE9VMRA67tZW5lhWdDDRwUFFPdHHA/dnNDXCadhEWr516PF9foY2u+7WNJSEj0gwCSgxIgXYmYVYeeVYueVYuRVYeeUo18g9TzQcpmjHudh6nznBHVSpQ8xyHefC8n4yppy+5tIb1tndlcX2rmzeas9hT3cWYR1UrorUfqaHcTwOTn7ktJzjcXjE/wjceAYHZUWafOvHdQyvge7V0Tv1cWnQ4I6eXstznOg0CYrc6PQJQ6dSiNxH1usAP3qcb4xhOQIacaMcRDqZDnY67ZLwEUlEAuo8p9CwMwsIZxdHblnF2NklWLllKDO+HZaBTYXWxnTtGNO1NqZpx5hOK9NVExcYXrJ0GzKA3NPv94Ogi9cCObwTyqTBzqBVd2HnKJyZDmqBQuUGRn9K7kROZEZZvUPHaDciodSuo/nOrWZkquGTvw1MDJfjDD6P3Bxyh8zRk+gRrHs1jW4t0qy6OzrUzsDQOp1aJIy6dZ2gBI9IIRJQaU5pBrY7FycjHzvTE7lleLAz87EzPISzCtF1nQL8lGs+SrRuSmiiTNvJZK2Dcq09et9BGZ0Y2vDQUES+IJsMg27doMuI/kI3dDp1gyZlcgQXXs2g29Tpy9BQ7hPf4fQDvp4ok0xCbSH0Lj1y69TRujT07mitKBoMGdFZQjNRZMRmEB2cRTSDyPOB2UWzVGQInaG37AmYCM6vabhKSjjU2REbTHTooKIDAeTTNMISPCINJU1ALVq0iBUrVmAYBhs3bmT9+vVx610uF2vWrGH27Nl0d3fzwAMP0NraCsDtt9/OsmXLsG2bRx99lDfeeGMiDiFhFBrKcEdvLpThxjEzccxMdNON223idrvIcBlkuXUy3QZ/2t7AZYuvJccMk0M/eQTwaD7ytBayNT/Zei9ZWoAsvZcMvZeQPjAnjE6/Hpk/JqDrNOka+zSNXl2nTyvEr+v06Dp+Xcen6/h1Df9JOkNqSmHYYNpgOJH7LBvyehWmL/LcFQbTAVdYRZ5Hl7lscIXAHQB3n4Y7ABkBDXc/uIMaGUGYO2s6h/buxY2NWxGbqtqtIkHk5tTdl8aTRSRw/LqGf8jUCb1aZPqEHl3DN+SxPzofz6q/u1dauYnzVlIElK7rrFy5knvuuYe2tjZqa2upq6vj0KFDsW2WLVtGT08Pt9xyCzU1Ndx1112sXr2aGTNmUFNTw6233kpJSQmPP/44f/VXf4XjnHpGVaezhb+ZEohbpnGSU0pn+OP05OPFAdF9aEM3VIA22K10aFOA+HJFt1Eqbr1yQPWB6ots6QDq6ODelBq41/ADvQp0BZoCTWWjO9lo0WWxm6PIUZDnDF02eG9Eb7qjMBw79jx2s8chHHbvZn6i9zFEGOImfxs6KVxA0+iNBrs/uqxXi4R3us3RI8R4SYqAqq6upqmpiebmZgC2bNnC4sWL4wLq6quv5mc/+xkAL7/8MitXrgRg8eLFbNmyBcuyaGlpoampierqanbt2nXKfdr6FAIljybmgMR5JzN6Kx7j921+EZYveGiM3zV5pPvxgRzjuUiKMyClpaW0tbXFnre1tVFaWjpsm2PHImOz2baN3+/H4/HELQfwer3DXjsi+TUrhBBJLSkCamKMTW95IYQQiZEUp/i8Xi9lZWWx52VlZXi93mHbTJo0Ca/Xi2EY5Obm0t3dHVs+oLS0dNhrR+LKd/jl9vvH7iCSzKpVq9L+4rrFjjOyAAAIzUlEQVQcY+pL9+MDOcbR+Di1Iy5PihpUQ0MDlZWVVFRUYJomNTU11NXVxW2zdetWli5dCsCSJUt4++23Aairq6OmpgaXy0VFRQWVlZXs2bNn3I9BCCHE2EqKGpRt26xdu5bHHnsMXdfZvHkzjY2N3HnnnTQ0NLB161Y2bdrEgw8+yIYNG/D5fKxevRqAxsZGXnrpJX79619j2zY/+MEPTtuCTwghRPJLioACqK+vp76+Pm7Zk08+GXscCoW4//6RT8mtW7eOdevWJbR8QgghxldSnOITQgghTiQBJYQQIilJQAkhhEhKElBCCCGSkgSUEEKIpCQBJYQQIilJQAkhhEhKElBCCCGSkgSUEEKIpCQBJYQQIilJQAkhhEhK2qJFi87LiZGef/55WlpaJroYQghx3quoqOCTn/zksOXnbUAJIYRIbnKKTwghRFKSgBJCCJGUJKCEEEIkJQkoIYQQSUkCSgghRFKSgBJCCJGUzIkuQCItWrSIFStWYBgGGzduZP369XHrXS4Xa9asYfbs2XR3d/PAAw/Q2to6QaU9O6c7xqVLl3L33Xfj9XoB+M1vfsOmTZsmoqhn5f777+eqq66is7OT5cuXj7jNvffey1VXXUV/fz///M//zL59+8a5lOfmdMd42WWX8fDDD9Pc3AzAK6+8Qm1t7XgX86yVlZWxZs0aioqKUErx+9//nmeeeWbYdqn8OY7mGFP5c3S73fzkJz/B5XJhGAYvv/wyTz31VNw2ifg+TduA0nWdlStXcs8999DW1kZtbS11dXUcOnQots2yZcvo6enhlltuoaamhrvuuovVq1dPXKHP0GiOEeCll15i7dq1E1PIc/SHP/yBDRs2sGbNmhHXX3nllVRWVnLLLbcwd+5cVq1axV//9V+PcynPzemOEWD79u3cd99941iqsWPbNo8//jj79u0jOzubn//857z55ptx/05T/XMczTFC6n6OoVCIu+++m76+PgzD4Kc//Sn19fXs3r07tk0ivk/T9hRfdXU1TU1NNDc3Ew6H2bJlC4sXL47b5uqrr+a5554D4OWXX+bDH/7wRBT1rI3mGFPdu+++i8/nO+n6xYsX8/zzzwOwe/ducnNzKS4uHq/ijYnTHWOqa29vj9WGAoEAhw4dorS0NG6bVP8cR3OMqa6vrw8A0zQxTROl4sd4SMT3adoGVGlpKW1tbbHnbW1tw/7BlJaWcuzYMSDyC8jv9+PxeMa1nOdiNMcIcO2117J+/XoeeughysrKxrOICTf0MwTwer1p98UAMG/ePH7xi1/wwx/+kAsuuGCii3PWysvLqaqqivvlDen1OZ7sGCG1P0dd11m3bh3PPfccb775Jnv27Ilbn4jv07QNKBGxdetWbr75Zr74xS/y1ltvpdQpTBGxd+9ebrrpJm6//XY2bNjA97///Yku0lnJysrie9/7Ho899hiBQGCii5MQpzrGVP8cHcfhjjvu4MYbb6S6upqZM2cmfJ9pG1BerzeutlBWVhZrKDB0m0mTJgFgGAa5ubl0d3ePaznPxWiO0efzYVkWABs3bmT27NnjWsZEG/oZQuRX3In/D1JdIBCInV6pr6/HNM2UqulD5O/ru9/9Li+++CKvvPLKsPXp8Dme7hjT4XME8Pv9bNu2jUWLFsUtT8T3adoGVENDA5WVlVRUVGCaJjU1NdTV1cVts3XrVpYuXQrAkiVLePvttyeiqGdtNMc49Dz+1VdfPeyibaqrq6uLjYI8d+5cent7aW9vn+BSja2ioqLY4+rqajRNS6kfUhBpqXj48GGefvrpEdenw+d4umNM5c+xoKCA3NxcADIyMli4cCGHDx+O2yYR36dpPZr5lVdeyYoVK9B1nc2bN7Nu3TruvPNOGhoa2Lp1K263mwcffJCqqip8Ph+rV6+ONQFNFac7xq9+9at87GMfw7ZtfD4fjzzyyLB/WMns29/+NpdffjkFBQV0dHTw1FNPYZqRxqfPPvssAPfddx8f+chHCAaDfOc73+G9996byCKfsdMd42c/+1luuukmbNsmGAzy+OOPs3Pnzgku9ejNnz+fn/70p+zfvx/HcQB44oknKC8vB9LjcxzNMaby5zhr1izWrFmDrutomsaf/vQnamtrE/59mtYBJYQQInWl7Sk+IYQQqU0CSgghRFKSgBJCCJGUJKCEEEIkJQkoIYQQSSltB4sVIp1lZGRw2223cd1111FeXk5XVxcvvPACTz75JLZtT3TxhBgTxtSpU7810YUQQoxecXExTzzxBNdccw3btm3jjTfeoLy8nCVLllBUVMRrr7020UUUYkxIPyghUohpmvzHf/wH06dP595772XHjh1AZAy4X/3qV5SVlXHDDTfQ0dExwSUV4tzJNSghUshtt93GnDlz+PGPfxwLJ4hMhfDqq69iGAaXXnrpBJZQiLEjASVEisjIyODWW2/F6/Xyu9/9btj6gXHdho75JkQqk4ASIkVcc8015Ofn88c//nHEhhButxuAcDg83kUTIiGkFZ8QKeKqq64CoKSkhK985SvD1i9cuBAgbuI/IVKZBJQQKWL+/PkAXHfddafcrrGxcTyKI0TCSUAJkQIyMzOpqKjgwIEDLF++fNj67OxsXnjhBdrb22ltbZ2AEgox9uQalBApoLS0FOCks8xeccUVuFwu6uvrx7NYQiSUBJQQKcDlcgFgWdaI66+//noANm/ePG5lEiLRJKCESAED058XFxcPWzd37lyuvPJKXn/9dfbs2TPeRRMiYSSghEgB3d3dNDY2cvHFFzNr1qzY8vLycv7pn/4Jv9/PI488MoElFGLsyVh8QqSInp4ePv7xj7NkyRJKSkq49tprWblyJW63m/vuu4+DBw9OdBGFGFMyFp8QKWTp0qXcdtttTJkyhe7ubl5//XVqa2tP2nhCiFQmASWEECIpyTUoIYQQSUkCSgghRFKSgBJCCJGUJKCEEEIkJQkoIYQQSUkCSgghRFKSgBJCCJGUJKCEEEIkJQkoIYQQSen/A/QbpGKWxbxzAAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x432 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VF7C72usbdze"
      },
      "source": [
        "We can now our simulation with the alignment energy alone."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "UDqvU3ZqbkIb",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "e55b6e6c-51f9-4443-a838-016d8c08db32"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  E_align = partial(align_fn, J_align=0.5, D_align=45., alpha=3.)\n",
        "  # Map the align energy over all pairs of boids. While both applications\n",
        "  # of vmap map over the displacement matrix, each acts on only one normal.\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  dR = space.map_product(displacement)(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  return 0.5 * np.sum(E_align(dR, N, N))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size, boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "W1M4PaLiw90c"
      },
      "source": [
        "Now the boids align with one another and already the simulation is displaying interesting behavior! "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "L2GHLnTarYLl"
      },
      "source": [
        "### Avoidance\n",
        "\n",
        "We can incorporate an avoidance rule that will keep the boids from bumping into one another. This will help them to form a flock with some volume rather than collapsing together. To this end, imagine a very simple model of boids that push away from one another if they get within a distance $D_{\\text{Avoid}}$ and otherwise don't repel. We can use a simple energy similar to Alignment but without any angular dependence,\n",
        "\n",
        "$$\n",
        "\\epsilon_{\\text{Avoid}}(\\Delta R_{ij}) = \\begin{cases}\\frac{J_{\\text{Avoid}}}{\\alpha}\\left(1 - \\frac{||\\Delta R_{ij}||}{D_{\\text{Avoid}}}\\right)^\\alpha & ||\\Delta R_{ij}||<D_{\\text{Avoid}} \\\\ 0 & \\text{otherwise}\\end{cases}\n",
        "$$\n",
        "\n",
        "This is implemented in the following Python function. Unlike the case of alignment, here we want boids to move away from one another and so we don't need a stop gradient on $\\Delta R$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "XSLebEhQw6fB"
      },
      "source": [
        "def avoid_fn(dR, J_avoid, D_avoid, alpha):\n",
        "  dr = space.distance(dR) / D_avoid\n",
        "  return np.where(dr < 1., \n",
        "                  J_avoid / alpha * (1 - dr) ** alpha, \n",
        "                  0.)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "0Yd8juOjyvPU"
      },
      "source": [
        "Plotting the energy we see that it is highest when boids are overlapping and then goes to zero smoothly until $||\\Delta R|| = D_{\\text{Align}}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "Njhwf_T33crC",
        "cellView": "form",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 441
        },
        "outputId": "faff9481-be60-452a-e657-31c34a0b62c9"
      },
      "source": [
        "#@title Avoidance Energy\n",
        "\n",
        "dr = np.linspace(0, 2., 60)\n",
        "dR = vmap(lambda r: np.array([0., r]))(dr)\n",
        "Es = vmap(partial(avoid_fn, J_avoid=1., D_avoid=1., alpha=3.))(dR)\n",
        "plt.plot(dr, Es, 'r', linewidth=3)\n",
        "\n",
        "plt.xlim([0, 2])\n",
        "\n",
        "format_plot('$r$', '$E$')\n",
        "finalize_plot()"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "image/png": "iVBORw0KGgoAAAANSUhEUgAAAagAAAGoCAYAAAATsnHAAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADh0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uMy4yLjIsIGh0dHA6Ly9tYXRwbG90bGliLm9yZy+WH4yJAAAgAElEQVR4nO3de3xU9YH38c9cciUXLuaGRPDGJbEW3QLRYlq6USu7VFGpVqUIhbYrtKjo2H0s+mypW+Ul22htn/ahDWapPizUutXYrhIVTdaUUlkp4BQwBiSEZEYIud9mMs8fDNMZrgmZmXMm832/Xn0xlzOTb6YDX8/v/M7vWIqKinyIiIiYjNXoACIiIqejghIREVNSQYmIiCmpoERExJRUUCIiYkoqKBERMSW70QGKioq4//77sdlsvPLKK6xfvz7k+blz53Lbbbfh9Xrp6uriySefZP/+/eTm5rJhwwYOHDgAwO7du1m9erURv4KIiESAoQVltVpZsWIFy5cvx+VyUVZWRlVVFfv37w9s8/rrr/Pyyy8DMHPmTJYvX84DDzwAQH19PQsWLDAiuoiIRJihQ3wFBQXU19fT0NCAx+OhsrKS4uLikG06OzsDt1NSUvD5dF6xiEg8MHQPKisrC5fLFbjvcrkoLCw8ZbvbbruNO++8k4SEBJYtWxZ4fOzYsZSXl9PR0cEvfvELduzYcdaf984779DR0RG+XyDC7HY7Ho/H6BgDEktZIbbyKmvkxFLeWMoKYLFYuOmmm4b0HoYfgxqIl156iZdeeokbbriBhQsXsmrVKo4cOcItt9xCa2srkyZN4qmnnuKuu+4K2eM6WUdHB7Nnz45i8qFxOBwxc1wtlrJCbOVV1siJpbyxlBWgrKxsyO9h6BCf2+0mOzs7cD87Oxu3233G7Tdv3hwYAuzr66O1tRWAPXv2cOjQIS666KLIBhYRkagxtKCcTif5+fnk5eVht9spKSmhqqoqZJtx48YFbn/+85/n4MGDAIwcORKr9Xj8sWPHkp+fT0NDQ/TCi4hIRBk6xOf1elmzZg2lpaVYrVYqKiqoq6tjyZIlOJ1Oqquruf3225k2bRoej4e2tjZWrVoFwNSpU1myZAkejwefz8fq1asDe1QiIhL7DD8GVVNTQ01NTchja9euDdwuLS097eu2bNnCli1bIhlNREQMpJUkRETElFRQIiJiSiooERExJRWUiIiYkgpKRERMSQUlIiKmpIISERFTUkGJiIgpqaBERMSU4qqgEtvbyejrMzqGiIgMQFwV1IimJlZ99JHRMUREZADiqqAApre0MCKGLvolIhKv4q6gAC7t6jI6goiInEN8FlQMXfZdRCRexWVBXXaWy8KLiIg5qKBERMSU4rKgLu3qwuLzGR1DRETOIq4Kqt9mA2CE10tuT4/BaURE5GziqqC8iYmB2xrmExExNxWUiIiYUnwVVFJS4LYKSkTE3OKroIL2oC5VQYmImFpcFVR/YiJe/+387m6SvN6zbi8iIsaJq4LyWSwcTE4Gjv/il2jJIxER04qrggL4KDU1cFvHoUREzCv+CmrEiMBtHYcSETGvuCuoWu1BiYjEhLgrqFOG+LTkkYiIKcVdQTUmJtLuX/Io0+MhS5eAFxExpbgrKCyWkL0oXRtKRMSc4q+g0HEoEZFYEJcFpanmIiLmp4JSQYmImFJcFtTHQQU1vrubhP5+A9OIiMjpxGVBddps1PtXNrf7fIzXkkciIqYTlwUFoRMlLtcwn4iI6cRtQek4lIiIuamg0Jp8IiJmpIJCe1AiImYUtwXVkJxMl/X4rz+mr49Rvb0GJxIRkWBxW1D9Fgsfp6QE7l+mmXwiIqZieEEVFRWxYcMGNm3axPz58095fu7cufz617+mvLycn//850yYMCHw3Ne//nU2bdrEhg0bmDFjxqB/tq4NJSJiXnYjf7jVamXFihUsX74cl8tFWVkZVVVV7N+/P7DN66+/zssvvwzAzJkzWb58OQ888AATJkygpKSEu+66iwsuuIBnn32WO+64g/5BnHQbchxKi8aKiJiKoXtQBQUF1NfX09DQgMfjobKykuLi4pBtOoP2bFJSUvD5r99UXFxMZWUlfX19HD58mPr6egoKCgb183UulIiIeRm6B5WVlYXL5Qrcd7lcFBYWnrLdbbfdxp133klCQgLLli0LvHbXrl2BbdxuN1lZWWf9eXa7HYfDEbif0NEB/mHFy3p7eeSBB/AlJAzpdwqnnJyckLxmFktZIbbyKmvkxFLeWMoaLoYW1EC99NJLvPTSS9xwww0sXLiQVatWndf7eDweVq9eHfLYVUlJXNjTg9Xj4eUf/pC9QceljOZwOE7Ja1axlBViK6+yRk4s5Y2lrABlZWVDfg9Dh/jcbjfZ2dmB+9nZ2bjd7jNuv3nz5sAQoNvtJicnJ/BcVlbWWV97JsGFNFHHoURETMPQgnI6neTn55OXl4fdbqekpISqqqqQbcaNGxe4/fnPf56DBw8CUFVVRUlJCQkJCeTl5ZGfn8+HH3446Ax7VFAiIqZk6BCf1+tlzZo1lJaWYrVaqaiooK6ujiVLluB0Oqmurub2229n2rRpeDwe2traAsN7dXV1vPnmm7z44ot4vV6efvrpQc3gO2Fv0ESJSSooERHTMPwYVE1NDTU1NSGPrV27NnC7tLT0jK8tLy+nvLx8SD8/eIjvss5OrD4f/RbLkN5TRESGzvATdY12JDGRT/0z91L7+xnX3W1wIhERARUUELoXpWE+ERFzUEERehxKEyVERMxBBcVJM/m0ooSIiCmooDjNEJ9/OSURETGOCgpoSEqizWYDINPjIUfXhhIRMZwKCsBi0YoSIiImo4Ly0wm7IiLmooLy26uJEiIipqKC8tOafCIi5qKC8vskJYUe/xJHOb29jOzrMziRiEh8U0H5eS2WkEvAay9KRMRYKqggGuYTETEPFVQQrcknImIeKqggWvJIRMQ8VFBBPk5NxeO/fVF3N6kez1m3FxGRyFFBBem1WtkfNFHicu1FiYgYRgV1kj2aySciYgoqqJNoRQkREXNQQZ1EM/lERMxBBXWS4EVjL+7qIrG/38A0IiLxSwV1kk67nYPJyQDYfT4u0TCfiIghVFCnsVcTJUREDKeCOg2dsCsiYjwV1GkET5SY0t5uYBIRkfilgjqNvwYV1GWdndg1UUJEJOpUUKfRkpBAfVISAIk+H5dpmE9EJOpUUGfgTEsL3NYwn4hI9KmgzsAZNMxXoJl8IiJRp4I6A+1BiYgYSwV1BntGjMDrvz2hq4tkr/es24uISHipoM6gy2Zjf0oKADa0Lp+ISLSpoM4iZJhPBSUiElUqqLNw6oRdERHDqKDOIngPqkAFJSISVSqos/goNZU+iwWAcT09ZHg8BicSEYkfKqiz6LNa2Re0svlk7UWJiESNCuocNFFCRMQYKqhzCFlRQntQIiJRo4I6B60oISJiDBXUOexPSaHTevxjyurrI6u31+BEIiLxwW50gKKiIu6//35sNhuvvPIK69evD3n+zjvv5Ctf+Qper5djx47xxBNP0NjYCEB1dTW1tbUANDU14XA4wp6v32Jhz4gRXNXWBhzfi3KPHh32nyMiIqEMLSir1cqKFStYvnw5LpeLsrIyqqqq2L9/f2CbvXv3snDhQnp6epg7dy5Lly5l5cqVAPT09LBgwYKI53SmpYUU1LsqKBGRiDN0iK+goID6+noaGhrweDxUVlZSXFwcss327dvp6ekBYPfu3WRnZ0c954fBK0poJp+ISFRYioqKfEb98FmzZlFUVMSPfvQjAL785S9TWFjImjVrTrv9ihUrOHLkCM8//zwAVVVV7Nu3D6/Xy/r163n33XfP+vPeeOMNKisrB50ztbGRG++7D4DetDReKy8H/wm8kZSTk0NTU1PEf044xFJWiK28yho5sZQ3lrICTJ48mUWLFg3pPQw/BjVQN954I5MnT+Y+f1EA3HrrrbjdbsaOHctzzz1HbW0thw4dOuN7eDweVq9ePfgf7vMxw25npMdDYns7L65aRX1y8vn8GoPicDjOL68BYikrxFZeZY2cWMobS1kBysrKhvwehg7xud3ukCG77Oxs3G73KdtNmzaNe++9F4fDQV9fX8jrARoaGti+fTsTJ06MTFCLhb9q4VgRkagytKCcTif5+fnk5eVht9spKSmhqqoqZJuJEyficDh4+OGHaW5uDjyenp5OQkICAJmZmVx55ZXU1dVFLqvOhxIRiSpDh/i8Xi9r1qyhtLQUq9VKRUUFdXV1LFmyBKfTSXV1NcuWLSM1NZUnnngC+Nt08gkTJvDII4/Q39+P1Wpl/fr1IbP/wk0TJUREosvwY1A1NTXU1NSEPLZ27drA7e9+97unfd3OnTu55557IpotWPAe1KSODmw+H94oTJQQEYlXWkligI4kJuJKTAQgpb+fCV1dBicSERneVFCDoIVjRUSiRwU1CLrCrohI9KigBmF3UEEVqqBERCJKBTUIH6al4fXfvqSzk1RdAl5EJGJUUIPQabNR678EvA0o1HRzEZGIUUEN0q6gYb4r/Cuci4hI+KmgBmlXenrg9hU6DiUiEjEqqEHaedIelMVn2GLwIiLDmgpqkOqTk2m2H1+AI8Pr5aLuboMTiYgMTyqowbJYQob5PqPjUCIiEaGCOg8hEyV0HEpEJCJUUOfh5ONQIiISfiqo8+BMS+PEKboXd3UxQifsioiEnQrqPHTbbHzkXzjWipY9EhGJBBXUeQo+DqWJEiIi4aeCOk87dcKuiEhEqaDO066TVjbXCbsiIuGlgjpPDUlJHE1IACDd69UVdkVEwkwFdb4sltDp5hrmExEJKxXUEAQfh7pSEyVERMJKBTUEuvSGiEjkqKCG4K9paXgsFgAmdHeToRN2RUTCRgU1BD1WK3v9V9gFKNBelIhI2Kighij4ONRnNFFCRCRsVFBDpBUlREQiQwU1RMHXhipob8eqE3ZFRMJCBTVEjYmJuP0n7I7o7+fizk6DE4mIDA8qqKGyWELPh9JxKBGRsFBBhYFO2BURCT8VVBh8EFRQV7W2go5DiYgMmQoqDPaNGEGH9fhHmdPbS25Pj8GJRERinwoqDLwnHYe6SsN8IiJDpoIKkw8yMgK3p7a2GphERGR4UEGFSfBxqKnagxIRGTIVVJg409Lo8S8ce1F3N2N6ew1OJCIS21RQYdJrtfJh0LJH2osSERkaFVQYBR+H+qyOQ4mIDIkKKoz+R8ehRETCRgUVRrvS0zlxycJLOzt1AUMRkSFQQYVRl83GnhEjgOMfrJY9EhE5f4YXVFFRERs2bGDTpk3Mnz//lOfvvPNOXnzxRdavX89PfvITcnNzA8/Nnj2bjRs3snHjRmbPnh3N2Gek86FERMLD0IKyWq2sWLGCBx98kK997Wtcf/31TJgwIWSbvXv3snDhQubPn89bb73F0qVLAcjIyGDRokUsXryYb3zjGyxatIj0oGNARlFBiYiEh6EFVVBQQH19PQ0NDXg8HiorKykuLg7ZZvv27fT417bbvXs32dnZAMyYMYNt27bR2tpKW1sb27Zto6ioKOq/w8l2pKfT7789qaODFK/X0DwiIrHKbuQPz8rKwuVyBe67XC4KCwvPuP2cOXOoqakJvLapqSnktVlZWWf9eXa7HYfDMcTU59b2wANkHjiAHfjBTTfhnjr1vN4nJycnKnnDIZayQmzlVdbIiaW8sZQ1XAwtqMG48cYbmTx5Mvfdd995v4fH42H16tVhTHWGn9PdzTz/7X2/+hX/Nz//vN7H4XBEJW84xFJWiK28yho5sZQ3lrIClJWVDfk9DB3ic7vdgSE7gOzsbNxu9ynbTZs2jXvvvReHw0FfX1/gtTk5Oed8rRF26DiUiMiQGVpQTqeT/Px88vLysNvtlJSUUFVVFbLNxIkTcTgcPPzwwzQ3Nwce37p1K9OnTyc9PZ309HSmT5/O1q1bo/0rnFbwwrEF7e0k9vefZWsRETkdQ4f4vF4va9asobS0FKvVSkVFBXV1dSxZsgSn00l1dTXLli0jNTWVJ554AoCmpiYcDgetra2sW7cusBtZVlZGq0n2Vo4kJvJJcjIXdXeT5PMxpb09ZK9KRETOzfBjUDU1NYGJDyesXbs2cPu73/3uGV9bUVFBRUVFxLINxQfp6VzU3Q0cX/ZIBSUiMjiGn6g7XOl8KBGRoVFBRUjwwrGfaWvD5vMZmEZEJPaooCKkMSmJxsREAEb09zOxo8PgRCIisUUFFSkWS8gw39Ua5hMRGRQVVAS9H1RQn2tpMTCJiEjsUUFF0J8zMwO3p7a1kaDzoUREBkwFFUGNSUnUJyUBkNzfzxXt7QYnEhGJHSqoCAvei9Iwn4jIwKmgIkwFJSJyflRQERY8UaKgo4NUXR9KRGRAVFARdiwhgb2pqQDYfT6tKiEiMkAqqCgIHuabpmE+EZEBUUFFwbbg86G0ByUiMiAqqCjYkZGBx2IB4LLOTkb5L7ooIiJnpoKKgi6bjV1paYH7f6dhPhGRc1JBRcmfNcwnIjIoKqgo2abzoUREBkUFFSUfpqXRaT3+cV/Y08NY/9V2RUTk9FRQUeKxWkMuv6G9KBGRs1NBRZGOQ4mIDJwKKoqCT9j9u5YWLLoMvIjIGamgouij1FSa7XYARns8XNLZaXAiERHzUkFFkc9iCVk8dpqG+UREzuisBbVgwQLGjx8frSxxQZffEBEZmLMW1De/+U3+/u//PuSxJP8VYuX8hFwGvrUVmy4DLyJyWoMe4rv77rt57bXXTvvcmDFjSElJGXKo4exQUhKHExMBGNHfT6EuAy8iclrndQwqM2gvINjNN9/MG2+8MaRAw57Fwp9GjgzcvebYMQPDiIiYV9gnSVj8q3bLmdWooEREzkmz+AywLejyG5M6Oxnd22twIhER81FBGaDTbucvQZffKNJelIjIKc5ZUD6tdhAR740aFbitYT4RkVPZz7XBvffey3XXXYfT6cTpdJKbmxuNXMPeH0eOZNknnwAwvaUFm8+HV8fvREQCzlpQ27ZtY9KkSYH/3XzzzYHnfvazn7Fv3z4++ugj9u3bR21tbcTDDie1KSm4EhPJ7u0lw+ulsL2dv6SnGx1LRMQ0zlpQ999/PwB5eXlMmTKFyZMnM2XKFCZOnMhnP/tZPvvZzwaGAPv7++nq6op84uHCYqFm5EhudrkAuKa5WQUlIhLknEN8AIcPH+bw4cO89dZbgcfGjRt3SmmlpaXpmNUghBTUsWP84qKLDE4kImIeAyqo06mvr6e+vp7NmzcHHhs/fjxTpkwJS7B4cGK6ud3nC0w3P+pfZUJEJN6FdZr5gQMH+K//+q9wvuWwpunmIiJnpvOgDKZVJURETk8FZbDggjox3VxERFRQhqtNTcXlP+50Yrq5iIiooIznn25+go5DiYgcZ3hBFRUVsWHDBjZt2sT8+fNPeX7q1Kk8//zzVFVVMWvWrJDnqqurKS8vp7y8nNWrV0crctjpOJSIyKnOe5p5OFitVlasWMHy5ctxuVyUlZVRVVXF/v37A9s0NjayatUq7r777lNe39PTw4IFC6KYODKCp5tP7ujQdHMREQzegyooKKC+vp6GhgY8Hg+VlZUUFxeHbNPY2EhtbS39w/jS6JpuLiJyKkP3oLKysnD5V1IAcLlcFBYWDvj1iYmJlJWV4fV6Wb9+Pe++++5Zt7fb7TgcjvPOG0mJv/0t/PrXACzMzeWKhx4iJyfHtHlPFktZIbbyKmvkxFLeWMoaLoYW1FDdeuutuN1uxo4dy3PPPUdtbS2HDh064/Yej8e0x6ou7ejg1/7bGVu3suapp1jxyCOmzXsyh8MRM1khtvIqa+TEUt5YygpQVlY25PcwdIjP7XaTnZ0duJ+dnY3b7R7U6wEaGhrYvn07EydODHvGaDl5uvln2toMTiQiYixDC8rpdJKfn09eXh52u52SkhKqqqoG9Nr09HQSEhIAyMzM5Morr6Suri6ScSPLYqEq6CKGxUePGhhGRMR4hg7xeb1e1qxZQ2lpKVarlYqKCurq6liyZAlOp5Pq6mqmTJnCk08+SXp6OjNnzmTx4sXcfffdTJgwgUceeYT+/n6sVivr168Pmf0Xi94dNYrbmpoAKG5u5k9aVUJE4pjhx6BqamqoqakJeWzt2rWB206nM+RCiSfs3LmTe+65J+L5oml7RgbtNhtpXi8X9vSQceCA0ZFERAxj+Im68jceq5X3gk7azfvTnwxMIyJiLBWUybwbdBxq7NatBiYRETGWCspkakaOpM9iAWBkXR25PT0GJxIRMYYKymQ67Xb+nJERuH+dZvOJSJxSQZnQu6NHB24XNzcbmERExDgqKBMKPh9qamsrGX19BqYRETGGCsqEjiQmstO/eKwd+LwWjxWROKSCMql3taqEiMQ5FZRJBR+HmtHSQtIwvtyIiMjpqKBM6pOUFNouvBCAlP5+prW0GJxIRCS6VFAm1jBjRuC2hvlEJN6ooEzs8PTpgdszm5uxavFYEYkjKigTa77sMtz+S4qM8nh0jSgRiSsqKDOzWqkOns2nk3ZFJI6ooEwueDbfF44eBQ3ziUicUEGZ3Pv+a0QBXNjTw6SODoMTiYhEhwrK5Pqs1pCTdq8/csTANCIi0aOCigGVY8YEbv/9kSNYNMwnInFABRUDtmVm0mK3A5Db28sV7e0GJxIRiTwVVAzwWK28HTRZokTDfCISB1RQMeLkYT6dtCsiw50KKkb8T0YGR/wn7Y7p6+Oq1laDE4mIRJYKKkb0Wyy8pWE+EYkjKqgYUnnBBYHbs44exaZLcIjIMKaCiiE709JoTEwEINPj0SU4RGRYU0HFEJ/FwltBkyV00q6IDGcqqBizOaigipubSdQwn4gMUyqoGPPXESOoT0oCIM3rpejYMYMTiYhEhgoq1lgsbA6aLKHZfCIyXKmgYtCbQcN8M5ubSfZ6DUwjIhIZKqgYVJuSwscpKQCk9PfzeV3IUESGIRVULLJYQpY+0mw+ERmOVFAxKrigrj12jMy+PgPTiIiEnwoqRh1MSWFXWhoACT6f9qJEZNhRQcWw32dlBW7/o8tlYBIRkfBTQcWwzWPG0GOxADCps5PLOjoMTiQiEj4qqBjWbrfzTtAK5//gdhuYRkQkvFRQMe61oGG+Gz/9FLuWPhKRYUIFFeP+nJlJk3+F81EeD5/X0kciMkyooGJcv8USMllitob5RGSYUEENA78PWpvv2uZmRvX2GphGRCQ8DC+ooqIiNmzYwKZNm5g/f/4pz0+dOpXnn3+eqqoqZs2aFfLc7Nmz2bhxIxs3bmT27NnRimw69SkpfJCeDoAd+PKnnxobSEQkDAwtKKvVyooVK3jwwQf52te+xvXXX8+ECRNCtmlsbGTVqlVs3rw55PGMjAwWLVrE4sWL+cY3vsGiRYtI9/8jHY+CJ0v8g9sNPp+BaUREhs7QgiooKKC+vp6GhgY8Hg+VlZUUFxeHbNPY2EhtbS39J81OmzFjBtu2baO1tZW2tja2bdtGUVFRNOObyltjxtBlPf5/56VdXUzWOVEiEuPsRv7wrKwsXEErILhcLgoLCwf82qamppDXZgXtRZyO3W7H4XCcX1gD5OTkDCqv+9lnuWjLFgBWjhvHjm9+M0LJTjXYrEaLpbzKGjmxlDeWsoaLoQUVbR6Ph9WrVxsdY8AcDseg8l7V0sLP/Lcv2LyZ0qNH6bVGZyd5sFmNFkt5lTVyYilvLGUFKCsrG/J7GDrE53a7yc7ODtzPzs7GPcBp0m63m5ycnPN67XD1QUYGDf7LwWd4vVyn60SJSAwztKCcTif5+fnk5eVht9spKSmhqqpqQK/dunUr06dPJz09nfT0dKZPn87WrVsjnNjcfBZLyGQJLSArIrHM0CE+r9fLmjVrKC0txWq1UlFRQV1dHUuWLMHpdFJdXc2UKVN48sknSU9PZ+bMmSxevJi7776b1tZW1q1bF9iNLCsro7W11chfxxR+n5XFN+rrsQJFLS2M6+qi3n/1XRGRWGL4MaiamhpqampCHlu7dm3gttPp5Oabbz7taysqKqioqIhovljTmJTEeyNHMtO/5NGtTU08e9LUfRGRWGD4iboSfi/l5gZu/6PbTbLXa2AaEZHzo4IahrZmZnIwORmAdK+XG7SyhIjEIBXUMOSzWPht0AzH25qatLKEiMQcFdQw9VpWFt3+c6AmdnZyZXu7wYlERAZHBTVMtdntvB60yvltjY0GphERGTwV1DD2UtAw36yjRxmty3CISAxRQQ1j+0aM4C9paQAk+HzcrBN3RSSGqKCGueAp57e4XNg0WUJEYoQKaph7e/RojiYkAJDd28t1R48anEhEZGBUUMNcn9XK74IW5L096BIlIiJmpoKKAy9nZ+Px3/671lYu7uw0NI+IyECooOKAOymJqtGjA/fnacq5iMQAFVSc+E3QlPPZbjejNOVcRExOBRUntmdk4BwxAoAkn4+vai9KRExOBRUvLBbWjx0buHtrUxOpHs9ZXiAiYiwVVBx5Z/RoPvGvcp7h9XKLTtwVERNTQcWRfouFF/LyAvfvPHyYhP5+AxOJiJyZCirO/FdWFp/6T9zN6uvjy7pWlIiYlAoqzvRarfxH0PJHdzc0YNHyRyJiQiqoOPRyTg7tNhsA47u7KW5uNjiRiMipVFBxqMNuD7ni7vxDh3TFXRExHRVUnNqYm0uPxQJAYUcHV7W2GpxIRCSUCipOHUlM5A9ZWYH78xsaDEwjInIqFVQceyEvjxOTzK9paeHyjg5D84iIBFNBxbH6lBTeDlpEdlF9vYFpRERCqaDiXPmFFwZuf7G5mSnt7QamERH5GxVUnNs3YgSVQXtR3zx40MA0IiJ/o4IS1ubn4/XfLmpp4aqWFkPziIiACkqAT1JSQmb0ffvgQZ0XJSKGU0EJAL8aN44+/3lRV7a3c+2xYwYnEpF4p4ISABqTkng5aHWJbx08qDX6RMRQKigJKB87li7r8a/ExM5OvnT0qMGJRCSeqaAk4GhiIhuDVjpfcvAgNu1FiYhBVFAS4oWxY2kLWun8Jrfb4EQiEq9UUBKizW7nhbFjA/e/UV+vq+6KiCFUUHKKjbm5HPVfdTe3t5evNjYanEhE4pEKSk7RZZYrPwsAABXZSURBVLOxLmgJpIX19Yzp7TUwkYjEIxWUnNbL2dl8nJICwIj+fpZ+8onBiUQk3qig5LS8Vis/njAhcP+mTz/lyrY24wKJSNxRQckZ/Tkzk7eCFpJdUVeHVdPORSRKVFByVj8ZP57uoJN3b3a5DE4kIvHC8IIqKipiw4YNbNq0ifnz55/yfEJCAqtWrWLTpk388pe/JNd/Imlubi5btmyhvLyc8vJyHA5HtKPHhcakJNYHTTv/1sGDZPT1GZhIROKFoQVltVpZsWIFDz74IF/72te4/vrrmRB03ANgzpw5tLW1MW/ePDZs2MDSpUsDz9XX17NgwQIWLFjA6tWro5w+frwwdiwNSUkAZHo8fFNX3hWRKDC0oAoKCqivr6ehoQGPx0NlZSXFxcUh21x33XX8/ve/B+Dtt9/mc5/7nBFR41qP1Urp+PGB+7c0NXF5R4eBiUQkHtiN/OFZWVm4go5puFwuCgsLT9mmqakJAK/XS3t7O5mZmQCMHTuW8vJyOjo6+MUvfsGOHTvO+vPsdntMDQXm5OSYJ6/PR9OqVeR88AE24N88Hqoefhj8l+gwVdYBiKW8yho5sZQ3lrKGi6EFNRRHjhzhlltuobW1lUmTJvHUU09x11130dnZecbXeDyemBoKdDgcpsp7kc/Hry0WEnw+LnA6+etDD/GK/xIdZst6LrGUV1kjJ5byxlJWgLKysiG/h6FDfG63m+zs7MD97Oxs3CctTup2u8nx/yNos9lIS0ujpaWFvr4+WltbAdizZw+HDh3ioosuil74OPRJSgobglY7/+4nn5Db02NgIhEZzgwtKKfTSX5+Pnl5edjtdkpKSqiqqgrZprq6mtmzZwMwa9Ys3n//fQBGjhyJ1T/9eezYseTn59PQ0BDdXyAO/So/nwPJyQCM8Hr5X7W1ujy8iESEoUN8Xq+XNWvWUFpaitVqpaKigrq6OpYsWYLT6aS6uppXX32Vxx9/nE2bNtHa2srKlSsBmDp1KkuWLMHj8eDz+Vi9enVgj0oip8dqZdWll/KL3buxAdNaW5mrc6NEJAIMPwZVU1NDTU1NyGNr164N3O7t7eXRRx895XVbtmxhy5YtkY4np7E7PZ0X8/KYf/gwAMsOHKDKP5FFRCRcDD9RV2LTL/PzqfMvJpva38/VP/0pFg31iUgYqaDkvPRarfzw0kvx+u9n7drFbdqLEpEwUkHJefswLY1fBy2DdN8nnzCuu9vARCIynKigZEh+NW4ctf6hvpT+fh6trdWK5yISFiooGZI+/6y+fv+U/6ltbSzWWn0iEgYqKBmyPWlp7Jk3L3B/4aFDXNvcbGAiERkOVFASFn+9/Xa2+tdIBHj8o4/I0/EoERkCFZSEh83G45ddRlNiIgAZXi//um8fif39BgcTkVilgpKwaUlI4NHLL6fPv8L55I4OHti/39hQIhKzVFASVrvT03k2+NpRLhezT1oAWERkIFRQEna/ycnhjTFjAvcdH3/MpbrAoYgMkgpKws9i4clLLgkshZTk8/Hk3r2M6uszOJiIxBIVlEREl83GP0+cSKf//KhxPT08/de/kuL1nuOVIiLHqaAkYg6kpPC/L7sssF5fQUcHP9y3D5tm9onIAKigJKKqRo9mzcUXB+5fe+wY36ur00UOReScVFAScS/n5LDuwgsD9//R7eabWg5JRM5BBSVR8X/HjePVrKzA/YWHDjG3sdHARCJidiooiQ6Lhacuvpj/Hjky8NCK/fv5wtGjBoYSETNTQUnUeK1Wvn/55eweMQIAG/DDffv40pEjxgYTEVNSQUlUddtsPDR5Mp8kJwNg9/n4wb59zHa5DE4mImajgpKoO5aQwLKCAvb7S8oGrPz4Y27TMSkRCaKCEkO4ExP5p8JC9qamBh57aP9+7m5oMDCViJiJCkoMc2JPamdaWuCxZZ98wpKDB3WelIiooMRYbXY790+ZwvsZGYHHFh06xMN1ddi14oRIXFNBieE6bTZWTJ7Me0FT0G91ufiJ08mo3l4Dk4mIkVRQYgo9ViuPTJwYcpmOqW1tPL9rFwXt7QYmExGjqKDENDxWK49fdhk/zc/nxOBedm8vP9u9m3/QNHSRuKOCEnOxWPj1hRfy4OTJtNpswPHrSX3/4495UMelROKKCkpMaevIkSz6zGeo9V/0EGBeUxO/2rWLy3V1XpG4oIIS0zqUnMySK67grdGjA49N7OykbNculhw8qL0pkWFOBSWm1mWz8ejll1M6fjw9FgtwfHmkRYcOsW7XLiZpAoXIsKWCEvOzWPiPvDzmX3klO9LTAw9f1tnJL3ft4tuffEKyLiUvMuyooCRmHExJ4b6CAn48fjzd1uNfXTuwoKGBjR98wM1NTdi0AoXIsKGCkpjSb7Gw0b839T9Be1NZfX18r66OF3bs4ItHjmipJJFhQAUlMak+OZmlBQX86OKLcSckBB4f393Nj/bt45e7dvF3LS0qKpEYpoKSmOWzWHglJ4d5U6fyf/LzafOfNwVQ2NHBc04n5Tt3MsflIknHqERijgpKYl6Pzca/X3ght0+dygt5eYHZfnB8Wvr/+vhjXtm+naUHDpDX3W1gUhEZDLvRAUTCpTUhgefGj2dTbi4LDh3ipk8/Jdl/rlSG18s9hw9z1+HD/HHkSJLefJOMvj5ag4YHRcRcVFAy7DQlJbH6kkv4PxddxD+6XNzW1MSFPT3A8SGDa48dg5/+lNeA7ZmZbBk9mi2jRtGcmGhobhEJpYKSYavNbuf/jR3Lf+Tlcc2xY9ze2EhRS0vgeTswvaWF6S0tPFRXh3PECHZkZPBBejo70tO1dyViMBWUDHv9Fgv/PWoU/z1qFHnd3cw6epS7k5MZvXdvYBsrxydWFHZ0cNfhwwB8nJLCjvR0/jpiBHWpqdSlpNBu118ZkWgx/G9bUVER999/PzabjVdeeYX169eHPJ+QkMBjjz3G5MmTaWlp4fvf/z6NjY0AfP3rX2fOnDl4vV5+/OMfs3XrViN+BYkhh5OTeXHsWMY5HKxbtYovNDfzpSNH+Gxb2ykzhi7p6uKSrq6Qx9wJCYGyakxKwp2YiDsxEVdCAp8mJuKxat6RSLgYWlBWq5UVK1awfPlyXC4XZWVlVFVVsX///sA2c+bMoa2tjXnz5lFSUsLSpUtZuXIlEyZMoKSkhLvuuosLLriAZ599ljvuuIN+LSAqA+ROSuI3ubn8JjeXDI+HK9va+GxrK1Pb2pjc0YH9NOdQZfX1keUfFjydZrudNruddpuN9hN/2mx02Gz0Wa30Wiz0Wa30Bf3Zb7HgA/o5PnXeB8f/55+NOO6dd7jh008j9jmEUyxlhdjKG0tZu8P0H2qGFlRBQQH19fU0NDQAUFlZSXFxcUhBXXfddfzqV78C4O2332bFihUAFBcXU1lZSV9fH4cPH6a+vp6CggJ27doV9d9DYl+r3U71qFFUjxoFQLLXS2F7O1e0t3NJZycXd3UxvquLxHOc+DvK42GUxxPecM88w7TwvmPkxFJWiK28sZQVcIbhPQwtqKysLFxBV0p1uVwUFhaesk1TUxMAXq+X9vZ2MjMzycrKCikjt9tNVlbWWX+e3W7H4XCE8TeIrJycnJjJG0tZYXB5O4BdwG6vl9SmJjIOHiTt0CFSjhwh5ehRkk/8eewYFu3Bi4SN4cegosnj8bB69WqjYwyYw+GImbyxlBXCmNdigQsugAsuwObzMbKvjzSvlzSP5/if/tupXi+JPh8J/f0k+Hwk+v+0+3xYfD6s8Lc//bdPKCgo4MMPPxx61iiIpawQW3ljKWuXzcakMLyPoQXldrvJzs4O3M/Ozsbtdp+yTU5ODm63G5vNRlpaGi0tLYHHT8jKyjrltSLR5LVYOJKYyJEwv6/jgQdipvxjKSvEVt5YygpQFob3MHTKkdPpJD8/n7y8POx2OyUlJVRVVYVsU11dzezZswGYNWsW77//PgBVVVWUlJSQkJBAXl4e+fn5MfNfFyIicm6G7kF5vV7WrFlDaWkpVquViooK6urqWLJkCU6nk+rqal599VUef/xxNm3aRGtrKytXrgSgrq6ON998kxdffBGv18vTTz+tGXwiIsOI4cegampqqKmpCXls7dq1gdu9vb08+uijp31teXk55eXlEc0nIiLG0FmFIiJiSiooERExJRWUiIiYkgpKRERMSQUlIiKmpIISERFTUkGJiIgpqaBERMSUVFAiImJKKigRETElFZSIiJiSpaio6OyXCB1G/vCHP3D48GGjY4iIDHt5eXncdNNNQ3qPuCooERGJHRriExERU1JBiYiIKamgRETElFRQIiJiSiooERExJRWUiIiYkt3oAOFSVFTE/fffj81m45VXXmH9+vUhzyckJPDYY48xefJkWlpa+P73v09jYyMAX//615kzZw5er5cf//jHbN261dCsd955J1/5ylfwer0cO3aMJ554IpC1urqa2tpaAJqamnA4HBHNOpC8s2fPZtmyZbjdbgB+85vf8Oqrrwaeu/feewF4/vnn+f3vf29o1uXLl3P11VcDkJyczKhRo7jhhhuA6H+2jz76KNdeey3Nzc3cc889p93mgQce4Nprr6W7u5tVq1axd+9eIPqf67my3nDDDcyfPx+LxUJnZyerV6/mo48+AuC3v/0tnZ2deL1evF4vixYtimjWgeS96qqrWL16NQ0NDQC88847lJWVAef+DkU769133x34jtpsNiZMmMDs2bNpbW2N+mebnZ3NY489xujRo/H5fPzud79j48aNp2wXru/tsCgoq9XKihUrWL58OS6Xi7KyMqqqqti/f39gmzlz5tDW1sa8efMoKSlh6dKlrFy5kgkTJlBSUsJdd93FBRdcwLPPPssdd9xBf3+/YVn37t3LwoUL6enpYe7cuYGsAD09PSxYsCAi2c43L8Cbb77JmjVrQh7LyMhg0aJFLFq0CJ/Px7p166iqqqKtrc2wrM8880zg9u23386kSZMC96P92b722mts2rSJxx577LTPX3PNNeTn5zNv3jwKCwtxOBwsXrw46p/rQLIePnyY++67j7a2NoqKivje977H4sWLA88vXbqUlpaWiOU72bnyAuzYsYOHHnoo5LGBft+jmfWFF17ghRdeAGDmzJnccccdtLa2Bp6P5mfr9Xp59tln2bt3L6mpqaxbt44//elPIZ9POL+3w2KIr6CggPr6ehoaGvB4PFRWVlJcXByyzXXXXRdo67fffpvPfe5zABQXF1NZWUlfXx+HDx+mvr6egoICQ7Nu376dnp4eAHbv3k12dnbE8pzLQPKeyYwZM9i2bRutra20tbWxbds2ioqKTJP1hhtu4I033ohYnnP54IMPQv6hOVlxcTF/+MMfgOPfg7S0NMaMGRP1z3UgWXfu3Bn4h8bo7yycO++ZDOX7fr4Gk/X6669n8+bNEc1zNkeOHAnsDXV2drJ//36ysrJCtgnn93ZYFFRWVhYulytw3+VynfKhZWVl0dTUBBz/r4D29nYyMzNDHgdwu92nvDbaWYPNmTOHmpqawP3ExETKyspYu3ZtxP/iwMDzfvGLX2T9+vU88cQTgX+cTv5sz/W7RisrQG5uLnl5ebz//vuBx6L92Z7Lmb6b0f5cB+vk76zP5+OZZ55h3bp13HzzzQYmC3XFFVfw7//+7/zbv/0bF198MTD4v5/RlJSURFFREVu2bAk8ZuRnm5uby8SJE9m9e3fI4+H83g6LIb7h6sYbb2Ty5Mncd999gcduvfVW3G43Y8eO5bnnnqO2tpZDhw4ZmPL4sZvNmzfT19fHLbfcwsqVK/nOd75jaKZzKSkp4e233w4ZyjXjZxtrrr76aubMmcO3vvWtwGPf/va3cbvdjBo1imeeeYYDBw7wwQcfGJgS9uzZw9y5c+nq6uKaa67hqaee4qtf/aqhmc5l5syZ/OUvfwnZ2zLqs01JSeFHP/oRpaWldHZ2RuznDIs9KLfbHTKkkJ2dHThgH7xNTk4OcPxAY1paGi0tLSGPw/H2P/m10c4KMG3aNO69914cDgd9fX0hrwdoaGhg+/btTJw4MWJZB5q3tbU1kPGVV15h8uTJgdcGf7Zn+l2jmfWE0w2VRPuzPZczfTej/bkO1KWXXso///M/43A4Qv4RPZGtubmZd955J6JD6APV2dlJV1cXADU1NdjtdjIzMwf1HYq2s31no/nZ2mw2/vVf/5XXX3+dd95555Tnw/m9HRYF5XQ6yc/PJy8vD7vdTklJCVVVVSHbVFdXM3v2bABmzZoVGNqpqqqipKSEhIQE8vLyyM/P58MPPzQ068SJE3E4HDz88MM0NzcHHk9PTychIQGAzMxMrrzySurq6iKWdaB5x4wZE7h93XXXBQ6Ybt26lenTp5Oenk56ejrTp0+P6AzJgWQFGD9+POnp6ezcuTPwmBGf7blUVVUFVoMuLCyko6ODI0eORP1zHYicnByefPJJfvCDH3Dw4MHA48nJyaSmpgZuz5gxg48//tiomAGjR48O3C4oKMBisdDS0jLg71C0jRgxgquuuop333038JhRn+2jjz7KgQMH2LBhw2mfD+f3dlgM8Xm9XtasWUNpaSlWq5WKigrq6upYsmQJTqeT6upqXn31VR5//HE2bdpEa2trYFZcXV0db775Ji+++CJer5enn346YjP4Bpp12bJlpKam8sQTTwB/m/I8YcIEHnnkEfr7+7Faraxfvz6is4sGmverX/0qM2fOxOv10trayg9/+EPg+J7VunXrAtN3y8rKzuvAdTizwvHhvZP/S9SIz/Zf/uVfuPrqqxk5ciS/+93v+OUvf4ndfvyv5Msvv8x7773Htddey6ZNm+jp6THscx1I1kWLFpGRkRGYFXdiyvPo0aN58skngeP/5f3GG2/wxz/+MaJZB5L3S1/6EnPnzsXr9dLT0xOYQXem75CRWQG+8IUvsHXrVrq7uwOvM+KzvfLKK7npppv46KOPKC8vB+DnP/85ubm5gbzh/N7qchsiImJKw2KIT0REhh8VlIiImJIKSkRETEkFJSIipqSCEhERU1JBiYiIKamgRETElFRQIiJiSiooERExJRWUiAldffXV1NTU8J3vfIeCggKeeuopXn/9dWpqagKXhhAZ7obFWnwiw82JK/1ecsklzJs3j/fee4///M//JCcnhwMHDhicTiQ6VFAiJnSioD7zmc/wT//0T6dcFE4kHmiIT8SEThRUaWmpyknilgpKxGSSk5PJz8/n6NGjvPbaa0bHETGMCkrEZC6//HJsNhvvvfcePp+uhiPxSwUlYjInhvc0tCfxTgUlYjInCsrpdBqcRMRYKigRk5k4cSK9vb3U1tYaHUXEUCooERNJSEjg4osvpra2Fo/HY3QcEUOpoERM5JJLLiEhIYE9e/YYHUXEcJaioiJNExIREdPRHpSIiJiSCkpERExJBSUiIqakghIREVNSQYmIiCmpoERExJRUUCIiYkoqKBERMSUVlIiImNL/BwLwcZVRGtSTAAAAAElFTkSuQmCC\n",
            "text/plain": [
              "<Figure size 432x432 with 1 Axes>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "G5E7-K61vQmd"
      },
      "source": [
        "We can now run a version of our simulation with both alignment and avoidance. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "YWxycV4NvVQ4",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "c4cb5b6f-c5be-4713-c749-955e62186903"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=1., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  # New Avoidance Code\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "  #\n",
        "\n",
        "  dR = space.map_product(displacement)(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  return 0.5 * np.sum(E_align(dR, N, N) + E_avoid(dR))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size,  boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "BPnoLPMXwDQp"
      },
      "source": [
        "The avoidance term in the energy stops the boids from collapsing on top of one another."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vdvp5lEQwHfB"
      },
      "source": [
        "### Cohesion\n",
        "\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2ry-azjayXIy"
      },
      "source": [
        "The final piece of Reynolds' boid model is cohesion. Notice that in the above simulation, the boids tend to move in the same direction but they also often drift apart. To make the boids behave more like schools of fish or birds, which maintain a more compact arrangement, we add a cohesion term to the energy. \n",
        "\n",
        "The goal of the cohesion term is to align boids towards the center of mass of their neighbors. Given a boid, $i$, we can compute the center of mass position of its neighbors as,\n",
        "\n",
        "$$\n",
        "\\Delta R_i = \\frac 1{|\\mathcal N|} \\sum_{j\\in\\mathcal N}\\Delta R_{ij}\n",
        "$$\n",
        "\n",
        "where we have let $\\mathcal N$ be the set of boids such that $||\\Delta R_{ij}|| < D_{\\text{Cohesion}}$. \n",
        "\n",
        "Given the center of mass displacements, we can define a reasonable cohesion energy as,\n",
        "\n",
        "$$\n",
        "\\epsilon_{Cohesion}\\left(\\widehat{\\Delta R}_i, N_i\\right) = \\frac12J_{\\text{Cohesion}}\\left(1 - \\widehat {\\Delta R}_i\\cdot N\\right)^2\n",
        "$$\n",
        "\n",
        "where $\\widehat{\\Delta R}_i = \\Delta R_i / ||\\Delta R_i||$ is the normalized vector pointing in the direction of the center of mass. This function is minimized when the boid is pointing in the direction of the center of  mass.\n",
        "\n",
        "We can implement the cohesion energy in the following python function. Note that as with alignment, we will have boids control their orientation and so we will insert a stop gradient on the displacement vector.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "j1vXoYDx6avU"
      },
      "source": [
        "def cohesion_fn(dR, N, J_cohesion, D_cohesion, eps=1e-7):\n",
        "  dR = lax.stop_gradient(dR)\n",
        "  dr = np.linalg.norm(dR, axis=-1, keepdims=True)\n",
        "  \n",
        "  mask = dr < D_cohesion\n",
        "\n",
        "  N_com = np.where(mask, 1.0, 0)\n",
        "  dR_com = np.where(mask, dR, 0)\n",
        "  dR_com = np.sum(dR_com, axis=1) / (np.sum(N_com, axis=1) + eps)\n",
        "  dR_com = dR_com / np.linalg.norm(dR_com + eps, axis=1, keepdims=True)\n",
        "  return f32(0.5) * J_cohesion * (1 - np.sum(dR_com * N, axis=1)) ** 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "oNqo8mJUDJ3y",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "728710a0-be23-4b54-fd1d-eb91c79abfa9"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=1., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "  \n",
        "  # New Cohesion Code\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.005, D_cohesion=40.)\n",
        "  #\n",
        "\n",
        "  dR = space.map_product(displacement)(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  return (0.5 * np.sum(E_align(dR, N, N) + E_avoid(dR)) + \n",
        "          np.sum(E_cohesion(dR, N)))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size,  boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kurQp4tK6FjT"
      },
      "source": [
        "Now the boids travel in tighter, more cohesive, packs. By tuning the range of the cohesive interaction and its strength you can change how strongly the boids attempt to stick together. However, if we raise it too high it can have some undesireable consequences."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gFapSAjV61q3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "8127c588-4880-4a2d-b10a-d69d535827d6"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=1., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.1, D_cohesion=40.)  # Raised to 0.05.\n",
        "\n",
        "  dR = space.map_product(displacement)(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  return (0.5 * np.sum(E_align(dR, N, N) + E_avoid(dR)) + \n",
        "          np.sum(E_cohesion(dR, N)))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size,  boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "8vB7NCrtEy8J"
      },
      "source": [
        "### Looking Ahead"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9Ojaz_QL7Uha"
      },
      "source": [
        "When the effect of cohesion is set to a large value, the boids cluster well. However, the motion of the individual flocks becomes less smooth and adopts an almost oscillatory behavior. This is caused by boids in the front of the pack getting pulled towards boids behind them.\n",
        "\n",
        "To improve this situation, we follow Reynolds and note that animals don't really look in all directions. The behavior of our flocks might look more realistic if we encorporated \"field of view\" for the boids. To this end, in both the alignment function and the cohesion function we will ignore boids that are outside of the line of sight for the boid. We will have a particularly simple definition for line of sight by first defining, $\\widehat{\\Delta R_{ij}} \\cdot N_i = \\cos\\theta_{ij}$ where $\\theta_{ij}$ is the angle between the orientation of the boid and the vector from the boid to its neighbor. \n",
        "\n",
        "Since most animals that display flocking behavior have eyes in the side of their head, as opposed to the front, we will define $\\theta_{\\text{min}}$ and $\\theta_{\\text{max}}$ to bound the angular field of view of the boids. Then, we assume each boid can see neighbors if $\\cos\\theta_{\\text{min}} < \\cos\\theta < \\cos\\theta_\\text{max}$."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4dXoJvOYA8vl"
      },
      "source": [
        "def field_of_view_mask(dR, N, theta_min, theta_max):\n",
        "  dr = space.distance(dR)\n",
        "  dR_hat = dR / dr\n",
        "  ctheta = np.dot(dR_hat, N)\n",
        "  # Cosine is monotonically decreasing on [0, pi].\n",
        "  return np.logical_and(ctheta > np.cos(theta_max),\n",
        "                        ctheta < np.cos(theta_min))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Dlyrm0UODHHY"
      },
      "source": [
        "We can then adapt the cohesion function to incorporate an arbitrary mask,"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "0BubUB7-C-Fz"
      },
      "source": [
        "def cohesion_fn(dR, N, mask, # New mask parameter.\n",
        "                J_cohesion, D_cohesion, eps=1e-7):\n",
        "  dR = lax.stop_gradient(dR)\n",
        "  dr = space.distance(dR)\n",
        "\n",
        "  mask = np.reshape(mask, mask.shape + (1,))\n",
        "  dr = np.reshape(dr, dr.shape + (1,))\n",
        "  \n",
        "  # Updated Masking Code\n",
        "  mask = np.logical_and(dr < D_cohesion, mask)\n",
        "  #\n",
        "  \n",
        "  N_com = np.where(mask, 1.0, 0)\n",
        "  dR_com = np.where(mask, dR, 0)\n",
        "  dR_com = np.sum(dR_com, axis=1) / (np.sum(N_com, axis=1) + eps)\n",
        "  dR_com = dR_com / np.linalg.norm(dR_com + eps, axis=1, keepdims=True)\n",
        "  return f32(0.5) * J_cohesion * (1 - np.sum(dR_com * N, axis=1)) ** 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "weLhL-GIDT1q"
      },
      "source": [
        "And finally run a simulation incorporating the field of view."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "hCw0HF1ADa5D",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "ddeb54ef-57bf-429b-eede-e991e92dfa6c"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=12., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.05, D_cohesion=40.)\n",
        "\n",
        "  dR = space.map_product(displacement)(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  # New FOV code.\n",
        "  fov = partial(field_of_view_mask, \n",
        "                theta_min=0.,\n",
        "                theta_max=np.pi / 3.)\n",
        "  # As before, we have to vmap twice over the displacement matrix, but only once\n",
        "  # over the normal.\n",
        "  fov = vmap(vmap(fov, (0, None)))\n",
        "  mask = fov(dR, N)\n",
        "  #\n",
        "\n",
        "  return (0.5 * np.sum(E_align(dR, N, N) * mask + E_avoid(dR)) + \n",
        "          np.sum(E_cohesion(dR, N, mask)))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size,  boids_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CUwj_eO_nt5U"
      },
      "source": [
        "## Extras\n",
        "\n",
        "Now that the core elements of the simulation are working well enough, we can add some extras fairly easily. In particular, we'll try to add some obstacles and some predators.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "6cQ7AHh3qLyd"
      },
      "source": [
        "\n",
        "### Obstacles\n",
        "\n",
        "The first thing we'll add are obstacles that the boids and (soon) the predators will try to avoid as they wander around the simulation. For the purposes of this notebook, we'll restrict ourselves to disk-like obstacles. Each disk will be described by a center position and a radius, $D_\\text{Obstacle}$. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "4VtFq3nqVsde"
      },
      "source": [
        "Obstacle = namedtuple('Obstacle', ['R', 'D'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Bl_t-_zDhQeb"
      },
      "source": [
        "Then we can instantiate some obstacles."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lt9RO3ychQDK"
      },
      "source": [
        "N_obstacle = 5\n",
        "\n",
        "R_rng, D_rng = random.split(random.PRNGKey(5))\n",
        "obstacles = Obstacle(\n",
        "    box_size * random.uniform(R_rng, (N_obstacle, 2)),\n",
        "    random.uniform(D_rng, (N_obstacle,), minval=30.0, maxval=100.0)\n",
        ")"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "WEGxQTl4V0fV"
      },
      "source": [
        "In a similar spirit to the energy functions above, we would like an energy function that encourages the boids to avoid obstacles. For this purpose we will pick an energy function that is similar in form to the alignment function above,\n",
        "\n",
        "$$\n",
        "\\epsilon_\\text{Obstacle}(\\Delta R_{io}, N_i, D_o) =  \\begin{cases}\\frac{J_\\text{Obstacle}}{\\alpha}\\left(1 - \\frac{\\|\\Delta R_{io}\\|}{D_o}\\right)^\\alpha\\left(1 + N_i\\cdot \\widehat{\\Delta R_{io}}\\right)^2 & \\|\\Delta R_{io}\\| < D_o \\\\ 0 & \\text{Otherwise}\\end{cases}\n",
        "$$\n",
        "\n",
        "for $\\Delta R_{io}$ the displacement vector between a boid $i$ and an obstacle $o$. This energy is zero when the boid and the obstacle are not overlapping. When they are overlapping, the energy is minimized when the boid is facing away from the obstacle.\n",
        "\n",
        "\\\n",
        "\n",
        "We can write down the boid-energy function in python.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sO5pgxSdeYWn"
      },
      "source": [
        "def obstacle_fn(dR, N, D, J_obstacle):\n",
        "  dr = space.distance(dR)\n",
        "  dR = dR / np.reshape(dr, dr.shape + (1,))\n",
        "  return np.where(dr < D,\n",
        "                  J_obstacle * (1 - dr / D) ** 2 * (1 + np.dot(N, dR)) ** 2,\n",
        "                  0.)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "dfr00bMWfS1O"
      },
      "source": [
        "Now we can run a simulation that includes obstacles."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "lzsn0Ca1fSg3",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "62a55ede-8d67-46c8-c206-e7b7c9e4a322"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  d = space.map_product(displacement)\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=12., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.05, D_cohesion=40.)\n",
        "\n",
        "  dR = d(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  fov = partial(field_of_view_mask, \n",
        "                theta_min=0.,\n",
        "                theta_max=np.pi / 3.)\n",
        "  fov = vmap(vmap(fov, (0, None)))\n",
        "  mask = fov(dR, N)\n",
        "\n",
        "  # New obstacle code\n",
        "  obstacles = state['obstacles']\n",
        "  dR_o = -d(boids.R, obstacles.R)\n",
        "  D = obstacles.D\n",
        "  E_obstacle = partial(obstacle_fn, J_obstacle=1000.)\n",
        "  E_obstacle = vmap(vmap(E_obstacle, (0, 0, None)), (0, None, 0))\n",
        "  #\n",
        "\n",
        "  return (0.5 * np.sum(E_align(dR, N, N) * mask + E_avoid(dR)) + \n",
        "          np.sum(E_cohesion(dR, N, mask)) + np.sum(E_obstacle(dR_o, N, D)))\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, speed=1.)\n",
        "\n",
        "boids_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids,\n",
        "    'obstacles': obstacles\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "\n",
        "display(render(box_size,  boids_buffer, obstacles))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QyCPMWIlOkHE"
      },
      "source": [
        "The boids are now successfully navigating obstacles in their environment."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qS_Y8LGBmI9R"
      },
      "source": [
        "### Predators\n",
        "\n",
        "Next we are going to introduce some predators into the environment for the boids to run away from. Much like the boids, the predators will be described by a position and an angle."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "rjSfZqe_Qh61"
      },
      "source": [
        "Predator = namedtuple('Predator', ['R', 'theta'])"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "gfYmTM4cd1sW"
      },
      "source": [
        "predators = Predator(R=np.array([[box_size / 2., box_size /2.]]),\n",
        "                     theta=np.array([0.0]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rKkjwZpSRMCx"
      },
      "source": [
        "The predators will also follow similar dynamics to the boids, swimming in whatever direction they are pointing at some speed that we can choose. Unlike in the previous versions of the simulation, predators naturally introduce some asymmetry to the system. In particular, we would like the boids to flee from the predators, but we want the predators to chase the boids. To achieve this behavior, we will consider a system reminiscient of a two-player game in which the boids move to minize an energy,\n",
        "\n",
        "$$\n",
        "E_\\text{Boid} = E_\\text{Align} + E_\\text{Avoid} + E_\\text{Cohesion} + E_\\text{Obstacle} + E_\\text{Boid-Predator}. \n",
        "$$\n",
        "\n",
        "Simultaneously, the predators move in an attempt to minimize a simpler energy,\n",
        "\n",
        "$$\n",
        "E_\\text{Predator} = E_\\text{Predator-Boid} + E_\\text{Obstacle}.\n",
        "$$\n",
        "\n",
        "To add predators to the environment we therefore need to add two rules, one that dictates the boids behavior near a predator and one for the behavior of predators near a group of boids. In both cases we will see that we can draw significant inspiration from behaviors that we've already developed.\n",
        "\n",
        "\\\n",
        "\n",
        "We will start with the boid-predator function since it is a bit simpler. In fact, we can use an energy that is virtually identical to the obstacle avoidance energy since the desired behavior is the same.\n",
        "\n",
        "$$\n",
        "\\epsilon_\\text{Boid-Predator}(\\Delta R_{ip}, N_i) = \\frac{J_\\text{Boid-Predator}}\\alpha\\left(1 - \\frac{\\|\\Delta R_{ip}\\|}{D_\\text{Boid-Predator}}\\right)^\\alpha (1 + \\widehat{\\Delta R_{ip}}\\cdot N_i)^2\n",
        "$$\n",
        "\n",
        "As before, this function is minimized when the boid is pointing away from the predators. Because we don't want the predators to experience this term we must include a stop-gradient on the predator positions.\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "HYnNci6CW0mN"
      },
      "source": [
        "def boid_predator_fn(R_boid, N_boid, R_predator, J, D, alpha):\n",
        "  N = N_boid\n",
        "  dR = displacement(lax.stop_gradient(R_predator), R_boid)\n",
        "  dr = np.linalg.norm(dR, keepdims=True)\n",
        "  dR_hat = dR / dr\n",
        "  return np.where(dr < D,\n",
        "                  J / alpha * (1 - dr / D) ** alpha * (1 + np.dot(dR_hat, N)),\n",
        "                  0.)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bHkPskM2joEJ"
      },
      "source": [
        "For the predator-boid function we can borrow the cohesion energy that we developed above to have predators that turn towards the center-of-mass of boids in their field of view."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "1NPpsNkhjlCI"
      },
      "source": [
        "def predator_boid_fn(R_predator, N_predator, R_boids, J, D, eps=1e-7):\n",
        "  # It is most convenient to define the predator_boid energy function\n",
        "  # for a single predator and a whole flock of boids. As such we expect shapes,\n",
        "  #   R_predator : (spatial_dim,)\n",
        "  #   N_predator : (spatial_dim,)\n",
        "  #   R_boids    : (n, spatial_dim,)\n",
        "  \n",
        "  N = N_predator \n",
        "\n",
        "  # As such, we need to vectorize over the boids.\n",
        "  d = vmap(displacement, (0, None))\n",
        "  dR = d(lax.stop_gradient(R_boids), R_predator)\n",
        "  dr = space.distance(dR)\n",
        "\n",
        "  fov = partial(field_of_view_mask, \n",
        "                theta_min=0.,\n",
        "                theta_max=np.pi / 3.)\n",
        "  # Here as well.\n",
        "  fov = vmap(fov, (0, None))\n",
        "\n",
        "  mask = np.logical_and(dr < D, fov(dR, N))\n",
        "  mask = mask[:, np.newaxis]\n",
        "\n",
        "  boid_count = np.where(mask, 1.0, 0)\n",
        "  dR_com = np.where(mask, dR, 0)\n",
        "  dR_com = np.sum(dR_com, axis=0) / (np.sum(boid_count, axis=0) + eps)\n",
        "  dR_com = dR_com / np.linalg.norm(dR_com + eps, keepdims=True)\n",
        "  return f32(0.5) * J * (1 - np.dot(dR_com, N)) ** 2"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pjzsy4mv08J_"
      },
      "source": [
        "Now we can modify our dynamics to also update predators."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "CcKoRMLM1KcG"
      },
      "source": [
        "def dynamics(energy_fn, dt, boid_speed, predator_speed):\n",
        "  # We extract common movement functionality into a `move` function.\n",
        "  def move(boids, dboids, speed):\n",
        "    R, theta, *_ = boids\n",
        "    dR, dtheta = dboids\n",
        "    n = normal(theta)\n",
        "\n",
        "    return (shift(R, dt * (speed * n + dR)), \n",
        "            theta + dt * dtheta)\n",
        "    \n",
        "  @jit\n",
        "  def update(_, state):\n",
        "    dstate = quantity.force(energy_fn)(state)\n",
        "\n",
        "    state['boids'] = Boids(*move(state['boids'], dstate['boids'], boid_speed))\n",
        "    state['predators'] = Predator(*move(state['predators'], \n",
        "                                        dstate['predators'], \n",
        "                                        predator_speed))\n",
        "\n",
        "    return state\n",
        "\n",
        "  return update"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "P5sY2ibN4VNM"
      },
      "source": [
        "Finally, we can put everything together and run the simulation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "fthzKiKI4U6Z",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 678
        },
        "outputId": "9e4c4f04-0ecd-4a21-d265-01ba0fa1fb2d"
      },
      "source": [
        "def energy_fn(state):\n",
        "  boids = state['boids']\n",
        "  d = space.map_product(displacement)\n",
        "  \n",
        "  E_align = partial(align_fn, J_align=12., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)), (0, 0, None))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.05, D_cohesion=40.)\n",
        "\n",
        "  dR = d(boids.R, boids.R)\n",
        "  N = normal(boids.theta)\n",
        "\n",
        "  fov = partial(field_of_view_mask, \n",
        "                theta_min=0.,\n",
        "                theta_max=np.pi / 3.)\n",
        "  fov = vmap(vmap(fov, (0, None)))\n",
        "  mask = fov(dR, N)\n",
        "\n",
        "  obstacles = state['obstacles']\n",
        "  dR_bo = -d(boids.R, obstacles.R)\n",
        "  D = obstacles.D\n",
        "  E_obstacle = partial(obstacle_fn, J_obstacle=1000.)\n",
        "  E_obstacle = vmap(vmap(E_obstacle, (0, 0, None)), (0, None, 0))\n",
        "  \n",
        "  # New predator code.\n",
        "  predators = state['predators']\n",
        "  E_boid_predator = partial(boid_predator_fn, J=256.0, D=75.0, alpha=3.)\n",
        "  E_boid_predator = vmap(vmap(E_boid_predator, (0, 0, None)), (None, None, 0))\n",
        "\n",
        "  N_predator = normal(predators.theta)\n",
        "  E_predator_boid = partial(predator_boid_fn, J=0.1, D=95.0)\n",
        "  E_predator_boid = vmap(E_predator_boid, (0, 0, None))\n",
        "\n",
        "  dR_po = -d(predators.R, obstacles.R)\n",
        "  #\n",
        "\n",
        "  E_boid = (0.5 * np.sum(E_align(dR, N, N) * mask + E_avoid(dR)) + \n",
        "            np.sum(E_cohesion(dR, N, mask)) + np.sum(E_obstacle(dR_bo, N, D)) + \n",
        "            np.sum(E_boid_predator(boids.R, N, predators.R)))\n",
        "  \n",
        "  E_predator = (np.sum(E_obstacle(dR_po, N_predator, D)) + \n",
        "                 np.sum(E_predator_boid(predators.R, N_predator, boids.R)))\n",
        "\n",
        "  return E_boid + E_predator\n",
        "\n",
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, boid_speed=1., predator_speed=.85)\n",
        " \n",
        "boids_buffer = []\n",
        "predators_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids,\n",
        "    'obstacles': obstacles,\n",
        "    'predators': predators\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "  predators_buffer += [state['predators']]\n",
        "\n",
        "display(render(box_size,  boids_buffer, obstacles, predators_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='400'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            400\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        },
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "<canvas id=\"canvas\"></canvas>\n",
              "<script>\n",
              "  Rg = null;\n",
              "  Ng = null;\n",
              "\n",
              "  var current_scene = {\n",
              "      R: null,\n",
              "      N: null,\n",
              "      is_loaded: false,\n",
              "      frame: 0,\n",
              "      frame_count: 0,\n",
              "      boid_vertex_count: 0,\n",
              "      boid_buffer: [],\n",
              "      predator_vertex_count: 0,\n",
              "      predator_buffer: [],\n",
              "      disk_vertex_count: 0,\n",
              "      disk_buffer: null,\n",
              "      box_size: 0\n",
              "  };\n",
              "\n",
              "  google.colab.output.setIframeHeight(0, true, {maxHeight: 5000});\n",
              "\n",
              "  async function load_simulation() {\n",
              "    buffer_size = 400;\n",
              "    max_frame = 800;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetObstacles', [], {});\n",
              "    data = result.data['application/json'];\n",
              "\n",
              "    if(data.hasOwnProperty('Disk')) {\n",
              "      current_scene = put_obstacle_disk(current_scene, data.Disk);\n",
              "    }\n",
              "\n",
              "    for (var i = 0 ; i < max_frame ; i += buffer_size) {\n",
              "      console.log(i);\n",
              "      result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetBoidStates', [i, i + buffer_size], {}); \n",
              "      \n",
              "      data = result.data['application/json'];\n",
              "      current_scene = put_boids(current_scene, data);\n",
              "    }\n",
              "    current_scene.is_loaded = true;\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "        'notebook.GetPredators', [], {}); \n",
              "    data = result.data['application/json'];\n",
              "    if (data.hasOwnProperty('R'))\n",
              "      current_scene = put_predators(current_scene, data);\n",
              "\n",
              "    result = await google.colab.kernel.invokeFunction(\n",
              "          'notebook.GetSimulationInfo', [], {});\n",
              "    current_scene.box_size = result.data['application/json'].box_size;\n",
              "  }\n",
              "\n",
              "  function initialize_gl() {\n",
              "    const canvas = document.getElementById(\"canvas\");\n",
              "    canvas.width = 640;\n",
              "    canvas.height = 640;\n",
              "\n",
              "    const gl = canvas.getContext(\"webgl2\");\n",
              "\n",
              "    if (!gl) {\n",
              "        alert('Unable to initialize WebGL.');\n",
              "        return;\n",
              "    }\n",
              "\n",
              "    gl.viewport(0, 0, gl.drawingBufferWidth, gl.drawingBufferHeight);\n",
              "    gl.clearColor(0.2, 0.2, 0.2, 1.0);\n",
              "    gl.enable(gl.DEPTH_TEST);\n",
              "\n",
              "    const shader_program = initialize_shader(\n",
              "        gl, VERTEX_SHADER_SOURCE_2D, FRAGMENT_SHADER_SOURCE_2D);\n",
              "    const shader = {\n",
              "      program: shader_program,\n",
              "      attribute: {\n",
              "          vertex_position: gl.getAttribLocation(shader_program, 'vertex_position'),\n",
              "      },\n",
              "      uniform: {\n",
              "          screen_position: gl.getUniformLocation(shader_program, 'screen_position'),\n",
              "          screen_size: gl.getUniformLocation(shader_program, 'screen_size'),\n",
              "          color: gl.getUniformLocation(shader_program, 'color'),\n",
              "      },\n",
              "    };\n",
              "    gl.useProgram(shader_program);\n",
              "\n",
              "    const half_width = 200.0;\n",
              "\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "    gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "\n",
              "    return {gl: gl, shader: shader};\n",
              "  }\n",
              "\n",
              "  var loops = 0;\n",
              "\n",
              "  function update_frame() {\n",
              "    gl.clear(gl.COLOR_BUFFER_BIT | gl.DEPTH_BUFFER_BIT);\n",
              "\n",
              "    if (!current_scene.is_loaded) {\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "      return;\n",
              "    }\n",
              "\n",
              "    var half_width = current_scene.box_size / 2.;\n",
              "    gl.uniform2f(shader.uniform.screen_position, half_width, half_width);\n",
              "    gl.uniform2f(shader.uniform.screen_size, half_width, half_width);\n",
              "\n",
              "    if (current_scene.frame >= current_scene.frame_count) {\n",
              "      if (!current_scene.is_loaded) {\n",
              "        window.requestAnimationFrame(update_frame);\n",
              "        return;\n",
              "      }\n",
              "      loops++;\n",
              "      current_scene.frame = 0;\n",
              "    }\n",
              "\n",
              "    gl.enableVertexAttribArray(shader.attribute.vertex_position);\n",
              "\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.boid_buffer[current_scene.frame]);\n",
              "    gl.uniform4f(shader.uniform.color, 0.0, 0.35, 1.0, 1.0);\n",
              "    gl.vertexAttribPointer(\n",
              "      shader.attribute.vertex_position,\n",
              "      2,\n",
              "      gl.FLOAT,\n",
              "      false,\n",
              "      0,\n",
              "      0\n",
              "    );\n",
              "    gl.drawArrays(gl.TRIANGLES, 0, current_scene.boid_vertex_count);\n",
              "\n",
              "    if(current_scene.predator_buffer.length > 0)  {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.predator_buffer[current_scene.frame]);\n",
              "      gl.uniform4f(shader.uniform.color, 1.0, 0.35, 0.35, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.predator_vertex_count);\n",
              "    }\n",
              "    \n",
              "    if(current_scene.disk_buffer) {\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, current_scene.disk_buffer);\n",
              "      gl.uniform4f(shader.uniform.color, 0.9, 0.9, 1.0, 1.0);\n",
              "      gl.vertexAttribPointer(\n",
              "        shader.attribute.vertex_position,\n",
              "        2,\n",
              "        gl.FLOAT,\n",
              "        false,\n",
              "        0,\n",
              "        0\n",
              "      );\n",
              "      gl.drawArrays(gl.TRIANGLES, 0, current_scene.disk_vertex_count);\n",
              "    }\n",
              "\n",
              "    current_scene.frame++;\n",
              "    if ((current_scene.frame_count > 1 && loops < 5) || \n",
              "        (current_scene.frame_count == 1 && loops < 240))\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "    \n",
              "    if (current_scene.frame_count > 1 && loops == 5 && current_scene.frame < current_scene.frame_count - 1)\n",
              "      window.requestAnimationFrame(update_frame);\n",
              "  }\n",
              "\n",
              "  function put_boids(scene, boids) {\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 8.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = i * boids + b;\n",
              "        var Nx = size * Math.cos(theta[ti]); //N[xi];\n",
              "        var Ny = size * Math.sin(theta[ti]); //N[yi];\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.boid_buffer.push(buffer);\n",
              "    }\n",
              "    scene.boid_vertex_count = boids * 3;\n",
              "    scene.frame_count += steps;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_predators(scene, boids) {\n",
              "    // TODO: Unify this with the put_boids function.\n",
              "    const R = decode(boids['R']);\n",
              "    const R_shape = boids['R_shape'];\n",
              "    const theta = decode(boids['theta']);\n",
              "    const theta_shape = boids['theta_shape'];\n",
              "\n",
              "    function index(i, b, xy) {\n",
              "      return i * R_shape[1] * R_shape[2] + b * R_shape[2] + xy; \n",
              "    }\n",
              "\n",
              "    var steps = R_shape[0];\n",
              "    var boids = R_shape[1];\n",
              "    var dimensions = R_shape[2];\n",
              "\n",
              "    if(dimensions != 2) {\n",
              "      alert('Can only deal with two-dimensional data.')\n",
              "    }\n",
              "\n",
              "    // First flatten the data.\n",
              "    var buffer_data = new Float32Array(boids * 6);\n",
              "    var size = 18.0;\n",
              "    for (var i = 0 ; i < steps ; i++) {\n",
              "      var buffer = gl.createBuffer();\n",
              "      for (var b = 0 ; b < boids ; b++) {\n",
              "        var xi = index(i, b, 0);\n",
              "        var yi = index(i, b, 1);\n",
              "        var ti = theta_shape[1] * i + b;\n",
              "        var Nx = size * Math.cos(theta[ti]);\n",
              "        var Ny = size * Math.sin(theta[ti]);\n",
              "        buffer_data.set([\n",
              "          R[xi] + Nx, R[yi] + Ny,\n",
              "          R[xi] - Nx - 0.5 * Ny, R[yi] - Ny + 0.5 * Nx,\n",
              "          R[xi] - Nx + 0.5 * Ny, R[yi] - Ny - 0.5 * Nx,             \n",
              "        ], b * 6);\n",
              "      }\n",
              "      gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "      gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "\n",
              "      scene.predator_buffer.push(buffer);\n",
              "    }\n",
              "    scene.predator_vertex_count = boids * 3;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  function put_obstacle_disk(scene, disk) {\n",
              "    const R = decode(disk.R);\n",
              "    const R_shape = disk.R_shape;\n",
              "    const radius = decode(disk.D);\n",
              "    const radius_shape = disk.D_shape;\n",
              "\n",
              "    const disk_count = R_shape[0];\n",
              "    const dimensions = R_shape[1];\n",
              "    if (dimensions != 2) {\n",
              "        alert('Can only handle two-dimensional data.');\n",
              "    }\n",
              "    if (radius_shape[0] != disk_count) {\n",
              "        alert('Inconsistent disk radius count found.');\n",
              "    }\n",
              "    const segments = 32;\n",
              "\n",
              "    function index(o, xy) {\n",
              "        return o * R_shape[1] + xy;\n",
              "    }\n",
              "\n",
              "    // TODO(schsam): Use index buffers here.\n",
              "    var buffer_data = new Float32Array(disk_count * segments * 6);\n",
              "    for (var i = 0 ; i < disk_count ; i++) {\n",
              "      var xi = index(i, 0);\n",
              "      var yi = index(i, 1);\n",
              "      for (var s = 0 ; s < segments ; s++) {\n",
              "        const th = 2 * s / segments * Math.PI;\n",
              "        const th_p = 2 * (s + 1) / segments * Math.PI;\n",
              "        const rad = radius[i] * 0.8;\n",
              "        buffer_data.set([\n",
              "          R[xi], R[yi],\n",
              "          R[xi] + rad * Math.cos(th), R[yi] + rad * Math.sin(th),\n",
              "          R[xi] + rad * Math.cos(th_p), R[yi] + rad * Math.sin(th_p),\n",
              "        ], i * segments * 6 + s * 6);\n",
              "      }\n",
              "    }\n",
              "    var buffer = gl.createBuffer();\n",
              "    gl.bindBuffer(gl.ARRAY_BUFFER, buffer);\n",
              "    gl.bufferData(gl.ARRAY_BUFFER, buffer_data, gl.STATIC_DRAW);\n",
              "    scene.disk_vertex_count = disk_count * segments * 3;\n",
              "    scene.disk_buffer = buffer;\n",
              "    return scene;\n",
              "  }\n",
              "\n",
              "  // SHADER CODE\n",
              "\n",
              "  const VERTEX_SHADER_SOURCE_2D = `\n",
              "    // Vertex Shader Program.\n",
              "    attribute vec2 vertex_position;\n",
              "    \n",
              "    uniform vec2 screen_position;\n",
              "    uniform vec2 screen_size;\n",
              "\n",
              "    void main() {\n",
              "      vec2 v = (vertex_position - screen_position) / screen_size;\n",
              "      gl_Position = vec4(v, 0.0, 1.0);\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  const FRAGMENT_SHADER_SOURCE_2D = `\n",
              "    precision mediump float;\n",
              "\n",
              "    uniform vec4 color;\n",
              "\n",
              "    void main() {\n",
              "      gl_FragColor = color;\n",
              "    }\n",
              "  `;\n",
              "\n",
              "  function initialize_shader(\n",
              "    gl, vertex_shader_source, fragment_shader_source) {\n",
              "\n",
              "    const vertex_shader = compile_shader(\n",
              "      gl, gl.VERTEX_SHADER, vertex_shader_source);\n",
              "    const fragment_shader = compile_shader(\n",
              "      gl, gl.FRAGMENT_SHADER, fragment_shader_source);\n",
              "\n",
              "    const shader_program = gl.createProgram();\n",
              "    gl.attachShader(shader_program, vertex_shader);\n",
              "    gl.attachShader(shader_program, fragment_shader);\n",
              "    gl.linkProgram(shader_program);\n",
              "\n",
              "    if (!gl.getProgramParameter(shader_program, gl.LINK_STATUS)) {\n",
              "      alert(\n",
              "        'Unable to initialize shader program: ' + \n",
              "        gl.getProgramInfoLog(shader_program)\n",
              "        );\n",
              "        return null;\n",
              "    }\n",
              "    return shader_program;\n",
              "  }\n",
              "\n",
              "  function compile_shader(gl, type, source) {\n",
              "    const shader = gl.createShader(type);\n",
              "    gl.shaderSource(shader, source);\n",
              "    gl.compileShader(shader);\n",
              "\n",
              "    if (!gl.getShaderParameter(shader, gl.COMPILE_STATUS)) {\n",
              "      alert('An error occured compiling shader: ' + gl.getShaderInfoLog(shader));\n",
              "      gl.deleteShader(shader);\n",
              "      return null;\n",
              "    }\n",
              "\n",
              "    return shader;\n",
              "  }\n",
              "\n",
              "  // SERIALIZATION UTILITIES\n",
              "  function decode(sBase64, nBlocksSize) {\n",
              "    var chrs = atob(atob(sBase64));\n",
              "    var array = new Uint8Array(new ArrayBuffer(chrs.length));\n",
              "\n",
              "    for(var i = 0 ; i < chrs.length ; i++) {\n",
              "      array[i] = chrs.charCodeAt(i);\n",
              "    }\n",
              "\n",
              "    return new Float32Array(array.buffer);\n",
              "  }\n",
              "\n",
              "  // RUN CELL\n",
              "\n",
              "  load_simulation();\n",
              "  gl_and_shader = initialize_gl();\n",
              "  var gl = gl_and_shader.gl;\n",
              "  var shader = gl_and_shader.shader;\n",
              "  update_frame();\n",
              "</script>\n"
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9e9MAZnXjtQe"
      },
      "source": [
        "We see that our predator now moves around chasing the boids."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rPVv0EsijR6H"
      },
      "source": [
        "### Internal State\n",
        "\n",
        "Until now, all of the data describing the boids, predators, and obstacles referred to their physical location and orientation. However, we can develop more interesting behavior if we allow the agents in our simulation to have extra data describing their internal state. As an example of this, we will allow predators to accelerate to chase boids if they get close.\n",
        "\n",
        "\\\n",
        "\n",
        "To this end, we add at extra piece of data to our predator, $t_\\text{sprint}$ which is the last time the predator accelerated. If the predator gets within $D_\\text{sprint}$ of a boid and it has been at least $T_\\text{sprint}$ units of time since it last sprinted it will accelerate. In practice, to accelerate the predator we will adjust its speed so that,\n",
        "\n",
        "$$\n",
        "s(t) = s_0 + s_1 e^{-(t - t_\\text{sprint}) / C\\tau_\\text{sprint}}\n",
        "$$\n",
        "\n",
        "where $s_0$ is the normal speed of the predator, $s_0 + s_1$ is the peak speed, and $\\tau_\\text{sprint}$ determines how long the sprint lasts for. In practice, rather than storign $t_\\text{sprint}$ we will record $\\Delta t = t - t_\\text{sprint}$ which is the time since the last sprint.\n",
        "\n",
        "\\\n",
        "\n",
        "Implementing this first requires that we add the necessary data to the predators."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "sG3oV_B6nVvJ"
      },
      "source": [
        "Predator = namedtuple('Predator', ['R', 'theta', 'dt']) "
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "-tQlq2k0pY8k"
      },
      "source": [
        "predators = Predator(R=np.array([[box_size / 2., box_size /2.]]),\n",
        "                     theta=np.array([0.0]),\n",
        "                     dt=np.array([0.]))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "zhRAmfz6nLQZ"
      },
      "source": [
        "def dynamics(energy_fn, dt, boid_speed, predator_speed):\n",
        "  # We extract common movement functionality into a `move` function.\n",
        "  def move(boids, dboids, speed):\n",
        "    R, theta, *_ = boids\n",
        "    dR, dtheta, *_ = dboids\n",
        "    n = normal(theta)\n",
        "\n",
        "    return (shift(R, dt * (speed * n + dR)), \n",
        "            theta + dt * dtheta)\n",
        "    \n",
        "  @jit\n",
        "  def update(_, state):\n",
        "    dstate = quantity.force(energy_fn)(state)\n",
        "\n",
        "    state['boids'] = Boids(*move(state['boids'], dstate['boids'], boid_speed))\n",
        "\n",
        "    # New code to accelerate the predators.\n",
        "    D_sprint = 65.\n",
        "    T_sprint = 300.\n",
        "    tau_sprint = 50.\n",
        "    sprint_speed = 2.0\n",
        "\n",
        "    # First we find the distance from each predator to the nearest boid.\n",
        "    d = space.map_product(space.metric(displacement))\n",
        "    predator = state['predators']\n",
        "    dr_min = np.min(d(state['boids'].R, predator.R), axis=1)\n",
        "\n",
        "    # Check whether there is a near enough boid to bother sprinting and if\n",
        "    # enough time has elapsed since the last sprint.\n",
        "    mask = np.logical_and(dr_min < D_sprint, predator.dt > T_sprint)\n",
        "    predator_dt = np.where(mask, 0., predator.dt + dt)\n",
        "\n",
        "    # Adjust the speed according to whether or not we're sprinting.\n",
        "    speed = predator_speed + sprint_speed * np.exp(-predator_dt / tau_sprint)\n",
        "\n",
        "    predator_R, predator_theta = move(state['predators'],\n",
        "                                      dstate['predators'], \n",
        "                                      speed)\n",
        "    state['predators'] = Predator(predator_R, predator_theta, predator_dt)\n",
        "    #\n",
        "\n",
        "    return state\n",
        "\n",
        "  return update"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "7xxQn6a7pW4Z",
        "colab": {
          "base_uri": "https://localhost:8080/",
          "height": 34
        },
        "outputId": "16ecbb3f-864c-41fd-8005-24ea3459483e"
      },
      "source": [
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, boid_speed=1., predator_speed=.85)\n",
        " \n",
        "boids_buffer = []\n",
        "predators_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids,\n",
        "    'obstacles': obstacles,\n",
        "    'predators': predators\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(400)):\n",
        "  state = lax.fori_loop(0, 50, update, state)\n",
        "  boids_buffer += [state['boids']]\n",
        "  predators_buffer += [state['predators']]\n",
        "\n",
        "display(render(box_size,  boids_buffer, obstacles, predators_buffer))"
      ],
      "execution_count": null,
      "outputs": [
        {
          "output_type": "display_data",
          "data": {
            "text/html": [
              "\n",
              "        <progress\n",
              "            value='301'\n",
              "            max='400',\n",
              "            style='width: 45%'\n",
              "        >\n",
              "            301\n",
              "        </progress>\n",
              "    "
            ],
            "text/plain": [
              "<IPython.core.display.HTML object>"
            ]
          },
          "metadata": {
            "tags": []
          }
        }
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "HZa6Zs9Xi--R"
      },
      "source": [
        "## Scaling Up"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LXU7JKCGlqHR"
      },
      "source": [
        "Up to this point, we have simulated a relatively small flock of $n = 200$ boids. In part we have done this because we compute, at each step, an $n\\times n$ matrix of distances. Therefore the computational complexity of the flocking simulation scales as $\\mathcal O(n^2)$. However, as Reynolds' notes we have built in a locality assumption so that no boids interact provided they are further apart than $D =\\max\\{D_{\\text{Align}}, D_\\text{Avoid}, D_\\text{Cohesion}\\}$. JAX MD provides tools to construct a set of candidates for each boid in about $\\mathcal O(n\\log n)$ time by predcomputing a list of neighbors for each boid. Using neighbor lists we can scale to much larger simulations.\n",
        "\n",
        "We create lists of all neighbors within a distance of $D + \\delta$ and pack them into an array of shape $n\\times n_\\text{max_neighbors}$. Using this technique we only need to rebuild the neighbor list if any particle has moved more than a distance of $\\delta$. We estimate `max_neighbors` from arrangements of particles and if any boid ever has more than this number of neighbors we must rebuild the neighbor list from scratch and recompile our simulation onto device.  \n",
        "\n",
        "To start with we setup a much larger system of boids."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZM28wDxD_h_S"
      },
      "source": [
        "# Simulation Parameters.\n",
        "box_size = 2400.0  # A float specifying the side-length of the box.\n",
        "boid_count = 2000  # An integer specifying the number of boids.\n",
        "obstacle_count = 10 # An integer specifying the number of obstacles.\n",
        "predator_count = 10 # An integer specifying the number of predators.\n",
        "dim = 2  # The spatial dimension in which we are simulating.\n",
        "\n",
        "# Create RNG state to draw random numbers.\n",
        "rng = random.PRNGKey(0)\n",
        "\n",
        "# Define periodic boundary conditions.\n",
        "displacement, shift = space.periodic(box_size)\n",
        "\n",
        "# Initialize the boids.\n",
        "# To generate normal vectors that are uniformly distributed on S^N note that\n",
        "# one can generate a random normal vector in R^N and then normalize it.\n",
        "rng, R_rng, theta_rng = random.split(rng, 3)\n",
        "\n",
        "boids = Boids(\n",
        "  R = box_size * random.uniform(R_rng, (boid_count, dim)),\n",
        "  theta = random.uniform(theta_rng, (boid_count,), maxval=2 * np.pi)\n",
        ")\n",
        "\n",
        "rng, R_rng, D_rng = random.split(rng, 3)\n",
        "obstacles = Obstacle(\n",
        "    R = box_size * random.uniform(R_rng, (obstacle_count, dim)),\n",
        "    D = random.uniform(D_rng, (obstacle_count,), minval=100, maxval=300.)    \n",
        ")\n",
        "\n",
        "rng, R_rng, theta_rng = random.split(rng, 3)\n",
        "predators = Predator(\n",
        "    R = box_size * random.uniform(R_rng, (predator_count, dim)),\n",
        "    theta = random.uniform(theta_rng, (predator_count,), maxval=2 * np.pi),\n",
        "    dt = np.zeros((predator_count,))\n",
        ")\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "jFuQCMw07_me"
      },
      "source": [
        "neighbor_fn = partition.neighbor_list(displacement,\n",
        "                                      box_size, \n",
        "                                      r_cutoff=45., \n",
        "                                      dr_threshold=10.,\n",
        "                                      capacity_multiplier=3)\n",
        "\n",
        "neighbors = neighbor_fn(boids.R)\n",
        "print(neighbors.idx.shape)"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mcWJUxYh_y0w"
      },
      "source": [
        "We see that dispite having 2000 boids, they each only have about 13 neighbors apiece at the start of the simulation. Of course this will grow over time and we will have to rebuild the neighbor list as it does. Next we make some minimal modifications to our energy function to rewrite the energy of our simulation to operate on neighbors. This mostly involves changing some of the vectorization patterns with `vmap` and creating a mask of which neighbors in the $n\\times n_\\text{max neighbors}$ arrays are filled. In JAX MD we use the pattern `mask = neighbors.idx == len(neighbors.idx)`. "
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "enwrFvl__Vtg"
      },
      "source": [
        "def energy_fn(state, neighbors):\n",
        "  boids = state['boids']\n",
        "  d = space.map_product(displacement)\n",
        "  \n",
        "  fov = partial(field_of_view_mask, \n",
        "                theta_min=0.,\n",
        "                theta_max=np.pi / 3.)\n",
        "  fov = vmap(vmap(fov, (0, None)))\n",
        "\n",
        "  E_align = partial(align_fn, J_align=12., D_align=45., alpha=3.)\n",
        "  E_align = vmap(vmap(E_align, (0, None, 0)))\n",
        "\n",
        "  E_avoid = partial(avoid_fn, J_avoid=25., D_avoid=30., alpha=3.)\n",
        "  E_avoid = vmap(vmap(E_avoid))\n",
        "\n",
        "  E_cohesion = partial(cohesion_fn, J_cohesion=0.05, D_cohesion=40.)\n",
        "\n",
        "  # New code to extract displacement vector to neighbors and normals.\n",
        "  R_neighbors = boids.R[neighbors.idx]\n",
        "  dR = -vmap(vmap(displacement, (None, 0)))(boids.R, R_neighbors)\n",
        "  N = normal(boids.theta)\n",
        "  N_neighbors = N[neighbors.idx]\n",
        "  #\n",
        "\n",
        "  # New code to add a mask over neighbors as well as field-of-view.\n",
        "  neighbor_mask = neighbors.idx < dR.shape[0]\n",
        "  fov_mask = np.logical_and(neighbor_mask, fov(dR, N))\n",
        "  #\n",
        "\n",
        "  obstacles = state['obstacles']\n",
        "  dR_bo = -d(boids.R, obstacles.R)\n",
        "  D = obstacles.D\n",
        "  E_obstacle = partial(obstacle_fn, J_obstacle=1000.)\n",
        "  E_obstacle = vmap(vmap(E_obstacle, (0, 0, None)), (0, None, 0))\n",
        "  \n",
        "  predators = state['predators']\n",
        "  E_boid_predator = partial(boid_predator_fn, J=256.0, D=75.0, alpha=3.)\n",
        "  E_boid_predator = vmap(vmap(E_boid_predator, (0, 0, None)), (None, None, 0))\n",
        "\n",
        "  N_predator = normal(predators.theta)\n",
        "  E_predator_boid = partial(predator_boid_fn, J=0.1, D=95.0)\n",
        "  E_predator_boid = vmap(E_predator_boid, (0, 0, None))\n",
        "\n",
        "  dR_po = -d(predators.R, obstacles.R)\n",
        "\n",
        "  E_boid = (0.5 * np.sum(E_align(dR, N, N_neighbors) * fov_mask + E_avoid(dR)) + \n",
        "            np.sum(E_cohesion(dR, N, fov_mask)) + np.sum(E_obstacle(dR_bo, N, D)) + \n",
        "            np.sum(E_boid_predator(boids.R, N, predators.R)))\n",
        "  \n",
        "  E_predator = (np.sum(E_obstacle(dR_po, N_predator, D)) + \n",
        "                 np.sum(E_predator_boid(predators.R, N_predator, boids.R)))\n",
        "\n",
        "  return E_boid + E_predator"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "lBbetwrhCsiu"
      },
      "source": [
        "Next we have to update our simulation to use and update the neighbor list."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "ZKjUyscEvN6D"
      },
      "source": [
        "def dynamics(energy_fn, dt, boid_speed, predator_speed):\n",
        "  # We extract common movement functionality into a `move` function.\n",
        "  def move(boids, dboids, speed):\n",
        "    R, theta, *_ = boids\n",
        "    dR, dtheta, *_ = dboids\n",
        "    n = normal(theta)\n",
        "\n",
        "    return (shift(R, dt * (speed * n + dR)), \n",
        "            theta + dt * dtheta)\n",
        "    \n",
        "  @jit\n",
        "  def update(_, state_and_neighbors):\n",
        "    state, neighbors = state_and_neighbors\n",
        "\n",
        "    #  New code to update neighbor list.\n",
        "    neighbors = neighbor_fn(state['boids'].R, neighbors) \n",
        "\n",
        "    dstate = quantity.force(energy_fn)(state, neighbors)\n",
        "    state['boids'] = Boids(*move(state['boids'], dstate['boids'], boid_speed))\n",
        "\n",
        "    # Predator acceleration.\n",
        "    D_sprint = 65.\n",
        "    T_sprint = 300.\n",
        "    tau_sprint = 50.\n",
        "    sprint_speed = 2.0\n",
        "\n",
        "    d = space.map_product(space.metric(displacement))\n",
        "    predator = state['predators']\n",
        "    dr_min = np.min(d(state['boids'].R, predator.R), axis=1)\n",
        "\n",
        "    mask = np.logical_and(dr_min < D_sprint, predator.dt > T_sprint)\n",
        "    predator_dt = np.where(mask, 0., predator.dt + dt)\n",
        "\n",
        "    speed = predator_speed + sprint_speed * np.exp(-predator_dt / tau_sprint)\n",
        "    speed = speed[:, np.newaxis]\n",
        "\n",
        "    predator_R, predator_theta = move(state['predators'],\n",
        "                                      dstate['predators'], \n",
        "                                      speed)\n",
        "    state['predators'] = Predator(predator_R, predator_theta, predator_dt)\n",
        "    #\n",
        "\n",
        "    return state, neighbors\n",
        "\n",
        "  return update"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "VgD9q_6iDrx2"
      },
      "source": [
        "And now we can conduct our larger simulation."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "b-LinFRy4Lkx"
      },
      "source": [
        "update = dynamics(energy_fn=energy_fn, dt=1e-1, boid_speed=1., predator_speed=.85)\n",
        "\n",
        "boids_buffer = []\n",
        "predators_buffer = []\n",
        "\n",
        "state = {\n",
        "    'boids': boids,\n",
        "    'obstacles': obstacles,\n",
        "    'predators': predators\n",
        "}\n",
        "\n",
        "for i in ProgressIter(range(800)):\n",
        "  new_state, neighbors = lax.fori_loop(0, 50, update, (state, neighbors))\n",
        "\n",
        "  # If the neighbor list can't fit in the allocation, rebuild it but bigger.\n",
        "  if neighbors.did_buffer_overflow:\n",
        "    print('REBUILDING')\n",
        "    neighbors = neighbor_fn(state['boids'].R)\n",
        "    state, neighbors = lax.fori_loop(0, 50, update, (state, neighbors))\n",
        "    assert not neighbors.did_buffer_overflow\n",
        "  else:\n",
        "    state = new_state\n",
        "\n",
        "  boids_buffer += [state['boids']]\n",
        "  predators_buffer += [state['predators']]\n",
        "\n",
        "display(render(box_size,  boids_buffer, obstacles, predators_buffer))"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GE1orUEXEnvT"
      },
      "source": [
        "At the end of the simulation we can see how large our neighbor list had to be to accomodate all of the boids."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "if1fbR3JEmPY"
      },
      "source": [
        "print(neighbors.idx.shape)"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}